diff --git a/isaaclab_arena/assets/background_library.py b/isaaclab_arena/assets/background_library.py index 4e5639bcfb..42def5c38f 100644 --- a/isaaclab_arena/assets/background_library.py +++ b/isaaclab_arena/assets/background_library.py @@ -13,6 +13,7 @@ from isaaclab_arena.assets.lightwheel_utils import acquire_lightwheel_asset from isaaclab_arena.assets.nucleus import ARENA_NUCLEUS_DIR from isaaclab_arena.assets.register import register_asset +from isaaclab_arena.relations.collision_mode import CollisionMode from isaaclab_arena.utils.pose import Pose @@ -189,6 +190,9 @@ def __init__(self, layout_id: int = 1, style_id: int = 1): )[0] ) super().__init__() + # PlaceableAsset.__init__ resets these; raw kitchen mesh keeps counter/floor concavities. + self.collision_mode = CollisionMode.MESH + self.repair_collision_mesh_non_watertight = False def get_viewer_cfg(self) -> ViewerCfg: # Looking in through the open front. diff --git a/isaaclab_arena/assets/object_reference.py b/isaaclab_arena/assets/object_reference.py index 9310b07c4d..e81dbb57f1 100644 --- a/isaaclab_arena/assets/object_reference.py +++ b/isaaclab_arena/assets/object_reference.py @@ -95,6 +95,11 @@ def get_world_bounding_box(self) -> AxisAlignedBoundingBox: quarters = quaternion_to_90_deg_z_quarters(parent_pose.rotation_xyzw) return box.rotated_90_around_z(quarters).translated(world_position) + def get_prim_path_in_parent_usd(self) -> str: + """Return the referenced prim's absolute path in its parent USD stage.""" + with open_stage(self.parent_asset.usd_path) as parent_stage: + return self.isaaclab_prim_path_to_original_prim_path(self.prim_path, self.parent_asset, parent_stage) + def get_collision_mesh(self) -> trimesh.Trimesh | None: """Return the referenced prim's collision mesh in its local frame, or None if unavailable.""" if not self._collision_mesh_loaded: diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index 66f6d1825e..b6c8f2f918 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -36,6 +36,7 @@ from isaaclab_arena.embodiments.embodiment_base import EmbodimentBase from isaaclab_arena.embodiments.franka.franka import franka_stack_events from isaaclab_arena.embodiments.robot_on_stand_utils import RobotPrimSpec, StandPrimSpec, compose_on_stand_usd +from isaaclab_arena.relations.collision_mode import CollisionMode from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.cameras import ArenaCameraCfg from isaaclab_arena.utils.pose import Pose @@ -67,8 +68,8 @@ class DroidEmbodimentBase(EmbodimentBase, ABC): which changes how far the stand extends below the root link. When manually placing the robot on floor, ``set_initial_pose`` z value and ``stand_height_m`` should be adjusted together to keep the bottom of stand fixed. - ``placement_bbox_stand_only`` uses the stand footprint for ``On`` / ``NextTo`` placement - instead of the full robot+stand USD bounds. + ``placement_bbox_stand_only`` uses the stand footprint for relation and collision placement + instead of the full robot+stand USD geometry. """ name = "droid" @@ -87,6 +88,8 @@ def __init__( super().__init__(enable_cameras, initial_pose, concatenate_observation_terms, arm_mode) self.stand_height_m = stand_height_m self.placement_bbox_stand_only = placement_bbox_stand_only + if placement_bbox_stand_only: + self.collision_mode = CollisionMode.BBOX self.scene_config = DroidSceneCfg() self.scene_config.robot.spawn.usd_path = compose_on_stand_usd( _DROID_ROBOT_PRIM, diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py index c8730b60ce..945b8b1304 100644 --- a/isaaclab_arena/environments/relation_solver_interface.py +++ b/isaaclab_arena/environments/relation_solver_interface.py @@ -11,7 +11,7 @@ from isaaclab_arena.relations.collision_mode import CollisionMode, get_object_collision_mode from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams -from isaaclab_arena.relations.placement_events import get_pose_from_layout, solve_and_place_objects +from isaaclab_arena.relations.placement_events import PlacementPoolHandle, get_pose_from_layout, solve_and_place_objects from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relations import get_anchor_objects from isaaclab_arena.utils.pose import PosePerEnv @@ -20,6 +20,7 @@ from isaaclab.managers import EventTermCfg from isaaclab_arena.assets.asset import Asset + from isaaclab_arena.assets.object_reference import ObjectReference from isaaclab_arena.assets.object_set import RigidObjectSet from isaaclab_arena.relations.collision_object import CollisionObject from isaaclab_arena.relations.placement_asset import PlaceableAsset @@ -29,11 +30,16 @@ def _get_passive_collision_objects( assets: Iterable[Asset | RigidObjectSet], include_background: bool = False, + background_mesh_exclusions: Iterable[ObjectReference] = (), ) -> list[CollisionObject]: """Load passive collision discovery only when relation placement needs it.""" from isaaclab_arena.relations.passive_collision_objects import get_passive_collision_objects - return get_passive_collision_objects(assets, include_background=include_background) + return get_passive_collision_objects( + assets, + include_background=include_background, + background_mesh_exclusions=background_mesh_exclusions, + ) def solve_and_apply_relation_placement( @@ -77,12 +83,18 @@ def solve_and_apply_relation_placement( # mutating the caller. placer_params.reachability_config = copy.copy(placer_params.reachability_config) if collision_objects is None and scene_assets is not None: + from isaaclab_arena.assets.object_reference import ObjectReference + scene_assets = list(scene_assets) + background_mesh_exclusions = [ + asset for asset in get_anchor_objects(assets) if isinstance(asset, ObjectReference) + ] collision_objects = _get_passive_collision_objects( scene_assets, include_background=_should_include_background_mesh( assets, scene_assets, placer_params.solver_params.collision_mode ), + background_mesh_exclusions=background_mesh_exclusions, ) placement_pool = PooledObjectPlacer( objects=assets, @@ -184,8 +196,7 @@ def _apply_dynamic_spawn_pose( func=solve_and_place_objects, mode="reset", params={ - "assets": assets, - "placement_pool": placement_pool, + "placement_pool": PlacementPoolHandle(placement_pool), }, ) diff --git a/isaaclab_arena/relations/background_collision_object.py b/isaaclab_arena/relations/background_collision_object.py index d2590766a6..51cfe180c8 100644 --- a/isaaclab_arena/relations/background_collision_object.py +++ b/isaaclab_arena/relations/background_collision_object.py @@ -67,15 +67,23 @@ def get_collision_mesh(self) -> trimesh.Trimesh: return self._mesh -def make_fixed_collision_objects(objects: Sequence[CollisionObject]) -> list[CollisionObject]: +def make_fixed_collision_objects( + objects: Sequence[CollisionObject], + excluded_prim_paths_by_object: dict[CollisionObject, Sequence[str]] | None = None, +) -> list[CollisionObject]: """Combine the objects' collision meshes into one FixedCollisionObject. - Objects in BBOX mode or without an extractable mesh are returned unchanged; - a whole-scene Background that cannot aggregate is an error. + Objects in BBOX mode or without an extractable mesh are returned unchanged. + Background mesh extraction failures are errors. + + Args: + objects: Fixed collision objects to aggregate. + excluded_prim_paths_by_object: USD prim subtrees omitted from individual objects' + extracted meshes, keyed by source object. """ from isaaclab_arena.assets.background import Background - mesh, skipped_objects = _combine_fixed_meshes(objects) + mesh, skipped_objects = _combine_fixed_meshes(objects, excluded_prim_paths_by_object) collision_objects: list[CollisionObject] = [] if mesh is not None: collision_objects.append(FixedCollisionObject(mesh)) @@ -104,18 +112,27 @@ def make_fixed_collision_objects(objects: Sequence[CollisionObject]) -> list[Col return collision_objects -def _combine_fixed_meshes(objects: Sequence[CollisionObject]) -> tuple[trimesh.Trimesh | None, list[CollisionObject]]: +def _combine_fixed_meshes( + objects: Sequence[CollisionObject], + excluded_prim_paths_by_object: dict[CollisionObject, Sequence[str]] | None = None, +) -> tuple[trimesh.Trimesh | None, list[CollisionObject]]: from isaaclab_arena.assets.background import Background from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache + from isaaclab_arena.utils.usd_helpers import AllCollisionMeshesExcludedError manager = WarpMeshAndSphereCache(device="cpu") + excluded_prim_paths_by_object = excluded_prim_paths_by_object or {} meshes = [] skipped_objects = [] for obj in objects: if obj.collision_mode == CollisionMode.BBOX: skipped_objects.append(obj) continue - mesh = manager.get_collision_mesh(obj) + excluded_prim_paths = excluded_prim_paths_by_object.get(obj, ()) + try: + mesh = manager.get_collision_mesh(obj, excluded_prim_paths=excluded_prim_paths) + except AllCollisionMeshesExcludedError: + continue if mesh is None: skipped_objects.append(obj) continue diff --git a/isaaclab_arena/relations/collision_mode.py b/isaaclab_arena/relations/collision_mode.py index d3efc9e72d..52d1438bc0 100644 --- a/isaaclab_arena/relations/collision_mode.py +++ b/isaaclab_arena/relations/collision_mode.py @@ -10,13 +10,15 @@ if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject + from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox class CollisionMode(Enum): """Collision-detection method for no-overlap constraints.""" BBOX = "bbox" - """Axis-aligned bounding box overlap volume (fast, conservative).""" + """Axis-aligned bounding box penetration (fast, conservative).""" MESH = "mesh" """Sphere-to-SDF queries against actual mesh geometry (accurate, slower).""" @@ -30,3 +32,32 @@ def get_object_collision_mode(obj: CollisionObject, default: CollisionMode) -> C def object_uses_mesh_collision(obj: CollisionObject, default: CollisionMode) -> bool: """Return True when the object's effective collision mode is MESH.""" return get_object_collision_mode(obj, default) == CollisionMode.MESH + + +def pair_is_covered_by_mesh_collision( + subject: CollisionObject, + obstacle: CollisionObject, + subject_bbox: AxisAlignedBoundingBox, + obstacle_bbox: AxisAlignedBoundingBox, + mesh_manager: WarpMeshAndSphereCache, + default_collision_mode: CollisionMode, + obstacle_is_fixed: bool, +) -> bool: + """Return whether the mesh loss represents a collision pair.""" + subject_has_mesh = ( + object_uses_mesh_collision(subject, default_collision_mode) + and mesh_manager.get_collision_mesh(subject) is not None + ) + obstacle_has_mesh = ( + object_uses_mesh_collision(obstacle, default_collision_mode) + and mesh_manager.get_collision_mesh(obstacle) is not None + ) + subject_has_mesh_or_invariant_bbox = subject_has_mesh or subject_bbox.is_batch_invariant() + if obstacle_is_fixed: + return obstacle_has_mesh and subject_has_mesh_or_invariant_bbox + obstacle_has_mesh_or_invariant_bbox = obstacle_has_mesh or obstacle_bbox.is_batch_invariant() + return ( + (subject_has_mesh or obstacle_has_mesh) + and subject_has_mesh_or_invariant_bbox + and obstacle_has_mesh_or_invariant_bbox + ) diff --git a/isaaclab_arena/relations/no_overlap_aabb.py b/isaaclab_arena/relations/no_overlap_aabb.py index a838c93457..a3bd89cd45 100644 --- a/isaaclab_arena/relations/no_overlap_aabb.py +++ b/isaaclab_arena/relations/no_overlap_aabb.py @@ -11,7 +11,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -from isaaclab_arena.relations.collision_mode import CollisionMode, object_uses_mesh_collision +from isaaclab_arena.relations.collision_mode import CollisionMode, pair_is_covered_by_mesh_collision from isaaclab_arena.relations.relation_loss_strategies import NoCollisionLossStrategy from isaaclab_arena.relations.relation_solver_state import RelationSolverState from isaaclab_arena.relations.relations import On @@ -26,20 +26,20 @@ class NoOverlapPair: """One directed overlap penalty: the subject box is pushed off the (detached) obstacle box. - Dimensions: B = batch_size (num envs). + B is the batch size. """ subject_min: torch.Tensor - """(B, 3) world-space min corner of the subject box.""" + """World-space subject minimum. Shape: [B, 3].""" subject_max: torch.Tensor - """(B, 3) world-space max corner of the subject box.""" + """World-space subject maximum. Shape: [B, 3].""" obstacle_min: torch.Tensor - """(B, 3) world-space min corner of the obstacle box.""" + """World-space obstacle minimum. Shape: [B, 3].""" obstacle_max: torch.Tensor - """(B, 3) world-space max corner of the obstacle box.""" + """World-space obstacle maximum. Shape: [B, 3].""" def compute_no_overlap_loss_aabb( @@ -163,13 +163,14 @@ def _fixed_pair_is_covered_by_mesh_collision( default_collision_mode: CollisionMode, ) -> bool: """Return True when MESH loss handles subject vs fixed obstacle.""" - obstacle_mesh = ( - mesh_manager.get_collision_mesh(obstacle) - if object_uses_mesh_collision(obstacle, default_collision_mode) - else None - ) - return obstacle_mesh is not None and _has_mesh_or_invariant_bbox( - state, subject, mesh_manager, default_collision_mode + return pair_is_covered_by_mesh_collision( + subject, + obstacle, + state.get_bbox(subject), + obstacle.get_bounding_box(), + mesh_manager, + default_collision_mode, + obstacle_is_fixed=True, ) @@ -181,23 +182,12 @@ def _dynamic_pair_is_covered_by_mesh_collision( default_collision_mode: CollisionMode, ) -> bool: """Return True when MESH loss handles a non-anchor object pair.""" - a_mesh = mesh_manager.get_collision_mesh(a) if object_uses_mesh_collision(a, default_collision_mode) else None - b_mesh = mesh_manager.get_collision_mesh(b) if object_uses_mesh_collision(b, default_collision_mode) else None - if a_mesh is None and b_mesh is None: - return False - return _has_mesh_or_invariant_bbox(state, a, mesh_manager, default_collision_mode) and _has_mesh_or_invariant_bbox( - state, b, mesh_manager, default_collision_mode + return pair_is_covered_by_mesh_collision( + a, + b, + state.get_bbox(a), + state.get_bbox(b), + mesh_manager, + default_collision_mode, + obstacle_is_fixed=False, ) - - -def _has_mesh_or_invariant_bbox( - state: RelationSolverState, - obj: PlaceableAsset, - mesh_manager: WarpMeshAndSphereCache, - default_collision_mode: CollisionMode, -) -> bool: - """Return True when MESH loss can represent obj as mesh or one bbox proxy.""" - mesh = mesh_manager.get_collision_mesh(obj) if object_uses_mesh_collision(obj, default_collision_mode) else None - if mesh is not None: - return True - return state.get_bbox(obj).is_batch_invariant() diff --git a/isaaclab_arena/relations/no_overlap_mesh.py b/isaaclab_arena/relations/no_overlap_mesh.py index 056e6f23f5..7eb94a86c3 100644 --- a/isaaclab_arena/relations/no_overlap_mesh.py +++ b/isaaclab_arena/relations/no_overlap_mesh.py @@ -154,12 +154,22 @@ def compute_no_overlap_loss_mesh( mesh_manager.warn_sdf_sentinel(sdf_values) sdf_values = clamp_sdf_sentinel(sdf_values) penetration = torch.relu(active_radii + clearance_m - sdf_values) - - pair_sum = torch.zeros(num_pairs, device=device, dtype=penetration.dtype) - pair_sum.index_add_(0, active_sphere_pair_id, penetration) - pair_mean = pair_sum / mesh_cache.pair_sphere_count + pair_subject_scale = ( + mesh_cache.pair_subject_bbox_max[:, b, :] - mesh_cache.pair_subject_bbox_min[:, b, :] + ).amax(dim=1) + assert torch.all(pair_subject_scale > 0), "Subject bounding boxes must have positive size." + normalized_penetration = penetration / pair_subject_scale[active_sphere_pair_id] + + pair_penetration = torch.zeros(num_pairs, device=device, dtype=penetration.dtype) + pair_penetration = pair_penetration.scatter_reduce( + 0, + active_sphere_pair_id, + normalized_penetration, + reduce="amax", + include_self=True, + ) active_pair_idx = active_pair.nonzero(as_tuple=True)[0] - total_loss[b] = total_loss[b] + slope * pair_mean[active_pair_idx].sum() + total_loss[b] = total_loss[b] + slope * pair_penetration[active_pair_idx].sum() if debug: print(f" [NoOverlap MESH] total_loss={total_loss.tolist()}") diff --git a/isaaclab_arena/relations/passive_collision_objects.py b/isaaclab_arena/relations/passive_collision_objects.py index f94d4ad71a..56d0e3cd36 100644 --- a/isaaclab_arena/relations/passive_collision_objects.py +++ b/isaaclab_arena/relations/passive_collision_objects.py @@ -23,7 +23,9 @@ def get_passive_collision_objects( - assets: Iterable[Asset | RigidObjectSet], include_background: bool = False + assets: Iterable[Asset | RigidObjectSet], + include_background: bool = False, + background_mesh_exclusions: Iterable[ObjectReference] = (), ) -> list[CollisionObject]: """Return relation-free scene assets that qualify as passive collision obstacles. @@ -34,6 +36,8 @@ def get_passive_collision_objects( assets: Scene assets to scan for relation-free fixed objects. include_background: If True, include Background assets and aggregate all mesh-capable objects into a single FixedCollisionObject. + background_mesh_exclusions: Object references whose USD subtrees are omitted + from an aggregated parent Background mesh. """ collision_objects: list[CollisionObject] = [] for asset in assets: @@ -76,5 +80,14 @@ def get_passive_collision_objects( ] if include_background: - return make_fixed_collision_objects(collision_objects) - return list(collision_objects) + excluded_prim_paths_by_object: dict[CollisionObject, list[str]] = {} + for reference in background_mesh_exclusions: + if reference.parent_asset in collision_object_set: + excluded_prim_paths_by_object.setdefault(reference.parent_asset, []).append( + reference.get_prim_path_in_parent_usd() + ) + return make_fixed_collision_objects( + collision_objects, + excluded_prim_paths_by_object=excluded_prim_paths_by_object, + ) + return collision_objects diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py index bed1ca3d7e..b2aa78de5b 100644 --- a/isaaclab_arena/relations/placement_events.py +++ b/isaaclab_arena/relations/placement_events.py @@ -26,6 +26,38 @@ PLACEMENT_RESET_EVENT_NAME = "placement_reset" +class PlacementPoolHandle: + """Opaque EventTermCfg param holding a runtime placement pool. + + Isaac Lab ``configclass._validate`` recursively walks any object with ``__dict__`` and has no + cycle guard. A live ``PooledObjectPlacer`` reaches placement assets (including embodiments with + cyclic scene configs) and overflows validation when stored directly in event params. + + This handle is the EventTermCfg-facing token: it intentionally has no ``__dict__`` so validation + stops here, while ``PooledObjectPlacer`` itself stays a normal class. Deep-copies share the same + pool instance (runtime state, not config). + """ + + __slots__ = ("pool",) + + def __init__(self, pool: PooledObjectPlacer) -> None: + self.pool = pool + + def __deepcopy__(self, memo: dict[int, object]) -> PlacementPoolHandle: + """Share the live pool across ``copy.deepcopy`` of EventTermCfg params.""" + memo[id(self)] = self + return self + + +def resolve_placement_pool(value: PooledObjectPlacer | PlacementPoolHandle | None) -> PooledObjectPlacer | None: + """Return the underlying pool, unwrapping a handle when present.""" + if value is None: + return None + if isinstance(value, PlacementPoolHandle): + return value.pool + return value + + def get_placement_pool(env) -> PooledObjectPlacer | None: """Return the pooled placer stored on the env reset event, or ``None`` when absent. @@ -39,7 +71,7 @@ def get_placement_pool(env) -> PooledObjectPlacer | None: term_cfg = env.unwrapped.event_manager.get_term_cfg(PLACEMENT_RESET_EVENT_NAME) except ValueError: return None - return term_cfg.params.get("placement_pool") + return resolve_placement_pool(term_cfg.params.get("placement_pool")) def get_rotation_xyzw(asset: PlaceableAsset) -> tuple[float, float, float, float]: @@ -113,8 +145,7 @@ def write_layout_to_sim( def solve_and_place_objects( env: ManagerBasedEnv, env_ids: torch.Tensor | None, - assets: list[PlaceableAsset], - placement_pool: PooledObjectPlacer, + placement_pool: PooledObjectPlacer | PlacementPoolHandle, ) -> None: """Coordinated reset event that draws layouts from the pool and writes poses. @@ -125,11 +156,14 @@ def solve_and_place_objects( Args: env: The Isaac Lab environment. env_ids: 1-D tensor of environment indices being reset. - assets: Assets participating in relation solving. - placement_pool: Runtime pool of solved placement layouts. + placement_pool: Runtime pool of solved placement layouts (or opaque handle). + Layout assets come from ``placement_pool.objects``. """ + placement_pool = resolve_placement_pool(placement_pool) + assert placement_pool is not None, "placement_reset event is missing its placement pool." if env_ids is None or len(env_ids) == 0: return + assets = placement_pool.objects reset_env_ids = env_ids.tolist() num_scene_envs = env.scene.env_origins.shape[0] assert ( diff --git a/isaaclab_arena/relations/placement_validators.py b/isaaclab_arena/relations/placement_validators.py index 65ee5612eb..a066cc1541 100644 --- a/isaaclab_arena/relations/placement_validators.py +++ b/isaaclab_arena/relations/placement_validators.py @@ -11,7 +11,12 @@ from collections.abc import Iterator from typing import TYPE_CHECKING, ClassVar, cast -from isaaclab_arena.relations.collision_mode import CollisionMode, get_object_collision_mode, object_uses_mesh_collision +from isaaclab_arena.relations.collision_mode import ( + CollisionMode, + get_object_collision_mode, + object_uses_mesh_collision, + pair_is_covered_by_mesh_collision, +) from isaaclab_arena.relations.placement_validation import PlacementCheck from isaaclab_arena.relations.placement_validator_registry import PlacementValidatorRegistry, register_validator from isaaclab_arena.relations.relation_loss_strategies import ( @@ -22,6 +27,7 @@ ) from isaaclab_arena.relations.relations import FaceTo, NextTo, NotNextTo, On, get_relation from isaaclab_arena.relations.warp_sdf_kernels import has_sdf_sentinel, mesh_sdf +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose from isaaclab_arena.utils.yaw import centers_in_target_frame, yaw_from_quat_xyzw, yaw_toward_positions @@ -30,7 +36,6 @@ from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.placement_asset import PlaceableAsset from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox class PlacementValidator(ABC): @@ -366,27 +371,57 @@ def validate_batch( bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: + batch_bboxes = self._stack_candidate_bboxes(bboxes) return [ - self._validate(positions[i], bboxes[i], orientations[i], collision_objects) for i in range(len(positions)) + self._validate( + positions[i], + bboxes[i], + orientations[i], + collision_objects, + batch_bboxes=batch_bboxes, + ) + for i in range(len(positions)) ] + @staticmethod + def _stack_candidate_bboxes( + bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], + ) -> dict[PlaceableAsset, AxisAlignedBoundingBox]: + """Stack candidate bounding boxes along the batch dimension.""" + return { + obj: AxisAlignedBoundingBox( + min_point=torch.cat([candidate[obj].min_point for candidate in bboxes]), + max_point=torch.cat([candidate[obj].max_point for candidate in bboxes]), + ) + for obj in bboxes[0] + } + def _validate( self, positions: dict[PlaceableAsset, tuple[float, float, float]], env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], orientations: dict[PlaceableAsset, float] | None, collision_objects: list[CollisionObject] | None, + batch_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox] | None = None, ) -> bool: """AABB overlap check, falling through to mesh penetration for mesh-collision objects.""" + batch_bboxes = batch_bboxes or env_bboxes use_mesh = self._should_validate_mesh(positions, collision_objects) no_overlap = self._validate_no_overlap( positions, env_bboxes, collision_objects=collision_objects, skip_mesh_pairs=use_mesh, + batch_bboxes=batch_bboxes, ) if no_overlap and use_mesh: - no_overlap = self._validate_no_overlap_mesh(positions, env_bboxes, orientations, collision_objects) + no_overlap = self._validate_no_overlap_mesh( + positions, + env_bboxes, + orientations, + collision_objects, + batch_bboxes=batch_bboxes, + ) return no_overlap def _should_validate_mesh( @@ -425,6 +460,7 @@ def _collect_skip_pairs( def _non_skip_pairs( self, positions: dict[PlaceableAsset, tuple[float, float, float]], + batch_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox] | None = None, skip_mesh_pairs: bool = False, ) -> Iterator[tuple[PlaceableAsset, PlaceableAsset]]: """Yield non-relation object pairs, optionally skipping pairs handled by mesh collision.""" @@ -439,17 +475,20 @@ def _non_skip_pairs( continue if (id(a), id(b)) in on_pairs: continue - if mesh_manager is not None and ( - ( - object_uses_mesh_collision(a, default_collision_mode) - and mesh_manager.get_collision_mesh(a) is not None - ) - or ( - object_uses_mesh_collision(b, default_collision_mode) - and mesh_manager.get_collision_mesh(b) is not None - ) - ): - continue + if id(a) in anchor_ids: + a, b = b, a + if mesh_manager is not None: + assert batch_bboxes is not None, "Mesh pair dispatch requires batched bounding boxes." + if pair_is_covered_by_mesh_collision( + a, + b, + batch_bboxes[a], + batch_bboxes[b], + mesh_manager, + default_collision_mode, + obstacle_is_fixed=b.is_anchor, + ): + continue yield a, b def _validate_no_overlap( @@ -458,14 +497,16 @@ def _validate_no_overlap( env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], collision_objects: list[CollisionObject] | None = None, skip_mesh_pairs: bool = False, + batch_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox] | None = None, ) -> bool: """AABB overlap check on pre-rotated env_bboxes. Skips On-pairs and anchor-anchor pairs.""" + batch_bboxes = batch_bboxes or env_bboxes clearance_m = self._params.solver_params.clearance_m margin = max(0.0, clearance_m - 1e-6) collision_objects = collision_objects or [] _, anchor_ids = self._collect_skip_pairs(positions) - for a, b in self._non_skip_pairs(positions, skip_mesh_pairs=skip_mesh_pairs): + for a, b in self._non_skip_pairs(positions, batch_bboxes, skip_mesh_pairs=skip_mesh_pairs): if self._pair_aabb_overlaps(env_bboxes[a], env_bboxes[b], positions[a], positions[b], 0.0, 0.0, margin): if self._params.verbose: print(f" Overlap between '{a.name}' and '{b.name}'") @@ -481,10 +522,14 @@ def _validate_no_overlap( continue obj_world = env_bboxes[obj].translated(positions[obj]) for background, background_world in background_worlds: - if ( - mesh_manager is not None - and object_uses_mesh_collision(background, default_collision_mode) - and mesh_manager.get_collision_mesh(background) is not None + if mesh_manager is not None and pair_is_covered_by_mesh_collision( + obj, + background, + batch_bboxes[obj], + background.get_bounding_box(), + mesh_manager, + default_collision_mode, + obstacle_is_fixed=True, ): continue if obj_world.overlaps(background_world, margin=margin).item(): @@ -510,30 +555,32 @@ def _validate_no_overlap_mesh( env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], orientations: dict[PlaceableAsset, float] | None = None, collision_objects: list[CollisionObject] | None = None, + batch_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox] | None = None, ) -> bool: """Sphere-to-SDF overlap check; both-meshless pairs fall back to AABB validation.""" + batch_bboxes = batch_bboxes or env_bboxes clearance_m = self._params.solver_params.clearance_m tolerance = max(0.0, clearance_m - 1e-6) mesh_manager = self._get_cpu_mesh_manager() mesh_manager.reset_sentinel_warning() - warned_no_mesh: set[str] = set() collision_objects = collision_objects or [] default_collision_mode = self._params.solver_params.collision_mode for a, b in self._non_skip_pairs(positions): + if not pair_is_covered_by_mesh_collision( + a, + b, + batch_bboxes[a], + batch_bboxes[b], + mesh_manager, + default_collision_mode, + obstacle_is_fixed=b.is_anchor, + ): + continue a_uses_mesh = object_uses_mesh_collision(a, default_collision_mode) b_uses_mesh = object_uses_mesh_collision(b, default_collision_mode) a_mesh = mesh_manager.get_collision_mesh(a) if a_uses_mesh else None b_mesh = mesh_manager.get_collision_mesh(b) if b_uses_mesh else None - if a_mesh is None and b_mesh is None: - for obj, uses_mesh, mesh in [(a, a_uses_mesh, a_mesh), (b, b_uses_mesh, b_mesh)]: - if uses_mesh and mesh is None and obj.name not in warned_no_mesh: - warned_no_mesh.add(obj.name) - print( - f" [NoCollision] MESH mode: '{obj.name}' has no collision mesh," - " falling back to AABB validation for this pair" - ) - continue a_pos = torch.tensor(positions[a], dtype=torch.float32) b_pos = torch.tensor(positions[b], dtype=torch.float32) @@ -581,13 +628,22 @@ def _validate_no_overlap_mesh( ) source_pos = torch.tensor(positions[source], dtype=torch.float32) for background in collision_objects: + if not pair_is_covered_by_mesh_collision( + source, + background, + batch_bboxes[source], + background.get_bounding_box(), + mesh_manager, + default_collision_mode, + obstacle_is_fixed=True, + ): + continue target_mesh = ( mesh_manager.get_collision_mesh(background) if object_uses_mesh_collision(background, default_collision_mode) else None ) - if target_mesh is None: - continue + assert target_mesh is not None, f"Mesh collision selected a meshless background '{background.name}'." target_pose = background.get_initial_pose() assert isinstance( target_pose, Pose diff --git a/isaaclab_arena/relations/relation_loss_strategies.py b/isaaclab_arena/relations/relation_loss_strategies.py index e61a700340..65d3939e56 100644 --- a/isaaclab_arena/relations/relation_loss_strategies.py +++ b/isaaclab_arena/relations/relation_loss_strategies.py @@ -12,7 +12,6 @@ from typing import TYPE_CHECKING from isaaclab_arena.relations.loss_primitives import ( - interval_overlap_axis_loss, linear_band_loss, single_boundary_linear_loss, single_point_linear_loss, @@ -458,15 +457,13 @@ def compute_loss( class NoCollisionLossStrategy: - """AABB no-overlap loss between object pairs.""" + """Size-normalized AABB penetration loss.""" + + def __init__(self, slope: float = 10.0): + """Initialize the loss. - def __init__( - self, - slope: float = 10.0, - ): - """ Args: - slope: Gradient magnitude for overlap loss. + slope: Weight for normalized penetration. """ self.slope = slope @@ -478,7 +475,7 @@ def compute_loss_batched( obstacle_min: torch.Tensor, obstacle_max: torch.Tensor, ) -> torch.Tensor: - """Overlap-volume no-overlap loss for boxes already reduced to world-space extents. + """Compute size-normalized penetration from world-space box extents. Args: clearance_m: Minimum clearance between boxes in meters. @@ -493,16 +490,24 @@ def compute_loss_batched( assert clearance_m >= 0, f"clearance_m must be non-negative, got {clearance_m}" obstacle_min = obstacle_min - clearance_m obstacle_max = obstacle_max + clearance_m - overlap_x = interval_overlap_axis_loss( - subject_min[..., 0], subject_max[..., 0], obstacle_min[..., 0], obstacle_max[..., 0] - ) - overlap_y = interval_overlap_axis_loss( - subject_min[..., 1], subject_max[..., 1], obstacle_min[..., 1], obstacle_max[..., 1] - ) - overlap_z = interval_overlap_axis_loss( - subject_min[..., 2], subject_max[..., 2], obstacle_min[..., 2], obstacle_max[..., 2] - ) - return self.slope * (overlap_x * overlap_y * overlap_z) + subject_center = (subject_min + subject_max) / 2 + obstacle_center = (obstacle_min + obstacle_max) / 2 + + # Pick the nearest separating face on each axis. Alternating pair tie-breaks keep the + # two directed passes of a coincident dynamic pair moving in opposite directions. + pair_index = torch.arange(subject_min.shape[0], device=subject_min.device) + tie_direction = torch.where(pair_index.remainder(2) == 0, -1.0, 1.0).view(-1, 1, 1) + center_delta = subject_center - obstacle_center + direction = torch.where(center_delta == 0, tie_direction, torch.sign(center_delta)) + depth_toward_negative = subject_max - obstacle_min + depth_toward_positive = obstacle_max - subject_min + axis_depth = torch.where(direction < 0, depth_toward_negative, depth_toward_positive) + + # A pair is collision-free if any axis is separated. + separation_depth = torch.relu(axis_depth.min(dim=-1).values) + subject_scale = (subject_max - subject_min).amax(dim=-1) + assert torch.all(subject_scale > 0), "Subject bounding boxes must have positive size." + return self.slope * separation_depth / subject_scale class AtPositionLossStrategy(UnaryRelationLossStrategy): diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py index de702a41b5..51ac245915 100644 --- a/isaaclab_arena/relations/relation_solver.py +++ b/isaaclab_arena/relations/relation_solver.py @@ -48,8 +48,7 @@ def __init__( params: Solver configuration parameters. If None, uses defaults. """ self.params = params or RelationSolverParams() - # High slope (vs 10-100 for relation strategies) so overlap avoidance dominates. - self._no_collision_strategy = NoCollisionLossStrategy(slope=10000.0) + self._no_collision_strategy = NoCollisionLossStrategy() self._last_loss_history: list[float] = [] self._last_position_history: list = [] self._last_loss_per_env: torch.Tensor | None = None diff --git a/isaaclab_arena/relations/warp_mesh_manager.py b/isaaclab_arena/relations/warp_mesh_manager.py index b8a16dc590..eff6fd9be1 100644 --- a/isaaclab_arena/relations/warp_mesh_manager.py +++ b/isaaclab_arena/relations/warp_mesh_manager.py @@ -11,6 +11,7 @@ import torch import trimesh from collections import defaultdict +from collections.abc import Sequence from heapq import heappop, heappush from typing import TYPE_CHECKING @@ -143,20 +144,42 @@ def warn_sdf_sentinel(self, sdf_values: torch.Tensor) -> None: "(no mesh face found). Collision detection may be incomplete for these points." ) - def get_collision_mesh(self, obj: CollisionObject) -> trimesh.Trimesh | None: + def get_collision_mesh( + self, + obj: CollisionObject, + excluded_prim_paths: Sequence[str] = (), + ) -> trimesh.Trimesh | None: """Return the cached collision mesh, extracting from USD on first access.""" from isaaclab_arena.assets.object import Object if not isinstance(obj, Object) or obj.usd_path is None: + assert not excluded_prim_paths, "USD prim exclusions require an Object with a usd_path." return obj.get_collision_mesh() usd_path = obj.usd_path - scale = tuple(obj.scale) - key = (usd_path, scale) + scale = obj.scale + exclusions = tuple(sorted(excluded_prim_paths)) + key = (usd_path, scale, exclusions) if key not in self._trimesh_cache: - from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd # deferred: pxr import + from isaaclab_arena.utils.usd_helpers import ( # deferred: pxr import + AllCollisionMeshesExcludedError, + NoCollisionMeshError, + UnsupportedCollisionGeometryError, + extract_trimesh_from_usd, + ) try: - self._trimesh_cache[key] = extract_trimesh_from_usd(usd_path, scale) + self._trimesh_cache[key] = extract_trimesh_from_usd( + usd_path, + scale, + excluded_prim_paths=exclusions, + ) + except AllCollisionMeshesExcludedError: + raise + except UnsupportedCollisionGeometryError as e: + print(f" [WarpMeshAndSphereCache] Could not extract mesh for '{obj.name}': {e}") + self._trimesh_cache[key] = None + except NoCollisionMeshError: + self._trimesh_cache[key] = None except ValueError as e: # Permanent: bad USD content, cache None to avoid re-parsing. print(f" [WarpMeshAndSphereCache] Could not extract mesh for '{obj.name}': {e}") diff --git a/isaaclab_arena/tests/test_embodiment_collision_mesh.py b/isaaclab_arena/tests/test_embodiment_collision_mesh.py index 8f5c746328..7c28b66132 100644 --- a/isaaclab_arena/tests/test_embodiment_collision_mesh.py +++ b/isaaclab_arena/tests/test_embodiment_collision_mesh.py @@ -12,9 +12,13 @@ def _test_embodiment_provides_robot_collision_mesh(simulation_app) -> bool: """Check the embodiment exposes its robot mesh so MESH mode does not fall back to the bbox proxy.""" from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment + from isaaclab_arena.relations.collision_mode import CollisionMode try: emb = DroidAbsoluteJointPositionEmbodiment() + stand_only_emb = DroidAbsoluteJointPositionEmbodiment(placement_bbox_stand_only=True) + assert emb.collision_mode is None + assert stand_only_emb.collision_mode == CollisionMode.BBOX mesh = emb.get_collision_mesh() assert mesh is not None, "embodiment must expose a collision mesh; None forces the loose bbox fallback" diff --git a/isaaclab_arena/tests/test_kitchen_bench_yaml_env.py b/isaaclab_arena/tests/test_kitchen_bench_yaml_env.py index df9775ae2a..0d3850b4c5 100644 --- a/isaaclab_arena/tests/test_kitchen_bench_yaml_env.py +++ b/isaaclab_arena/tests/test_kitchen_bench_yaml_env.py @@ -46,6 +46,45 @@ def _test_kitchen_bench_yaml_env_bringup(simulation_app, *, yaml_path: Path) -> return True +def _test_droid_stand_survives_pool_refill(simulation_app, *, yaml_path: Path) -> bool: + """Keep the instanceable Droid stand loaded across placement-pool refill.""" + from pxr import Usd, UsdGeom + + from isaaclab_arena.cli.isaaclab_arena_cli import arena_env_builder_cfg_from_argparse, get_isaaclab_arena_cli_parser + from isaaclab_arena.environment_spec.arena_env_graph_spec import ArenaEnvGraphSpec + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.relations.placement_events import get_placement_pool + + spec = ArenaEnvGraphSpec.from_yaml(yaml_path) + arena_env = spec.to_arena_env() + args_cli = get_isaaclab_arena_cli_parser().parse_args(["--num_envs", "1"]) + env = ArenaEnvBuilder(arena_env, arena_env_builder_cfg_from_argparse(args_cli)).make_registered() + + def visible_stand_meshes() -> tuple[str, ...]: + stand = env.unwrapped.scene.stage.GetPrimAtPath("/World/envs/env_0/Robot/panda_link0/stand_instanceable") + assert stand.IsValid() and stand.IsActive() and stand.IsLoaded() + meshes = tuple(prim for prim in Usd.PrimRange(stand, Usd.TraverseInstanceProxies()) if prim.IsA(UsdGeom.Mesh)) + assert meshes + assert all( + mesh.IsLoaded() and UsdGeom.Imageable(mesh).ComputeVisibility() != UsdGeom.Tokens.invisible + for mesh in meshes + ) + return tuple(str(mesh.GetPath()) for mesh in meshes) + + try: + env.reset() + mesh_paths = visible_stand_meshes() + placement_pool = get_placement_pool(env) + assert placement_pool is not None + placement_pool.sample_without_replacement(placement_pool.total_remaining) + env.reset() + assert visible_stand_meshes() == mesh_paths + finally: + env.close() + + return True + + @pytest.mark.parametrize( "yaml_path", _KITCHEN_BENCH_YAMLS, @@ -62,6 +101,16 @@ def test_kitchen_bench_yaml_env_bringup(yaml_path: Path): assert result, f"kitchen_bench bring-up failed for {yaml_path.name}" +def test_droid_stand_survives_pool_refill(): + """The Lightwheel kitchen keeps the instanceable Droid stand visible.""" + yaml_path = _KITCHEN_BENCH_DIR / "droid_pick_and_place_lightwheel_kitchen.yaml" + assert run_simulation_app_function( + _test_droid_stand_survives_pool_refill, + headless=True, + yaml_path=yaml_path, + ) + + if __name__ == "__main__": for path in _KITCHEN_BENCH_YAMLS: test_kitchen_bench_yaml_env_bringup(path) diff --git a/isaaclab_arena/tests/test_mesh_collision.py b/isaaclab_arena/tests/test_mesh_collision.py index 11b1e4cdde..0c39f4a77d 100644 --- a/isaaclab_arena/tests/test_mesh_collision.py +++ b/isaaclab_arena/tests/test_mesh_collision.py @@ -716,10 +716,66 @@ def test_mixed_mesh_aabb_varying_proxy_uses_aabb_fallback(): solver.solve([table, source, target], initial, env_bboxes=env_bboxes) losses = solver.last_loss_per_env - assert losses[1].item() > losses[0].item() + assert torch.all(losses > 0) + assert not torch.isclose(losses[1], losses[0]) assert solver._last_no_overlap_pair_count > 0 +@requires_warp +def test_validator_uses_batched_bboxes_for_mesh_dispatch(monkeypatch): + """Candidate-varying proxies use AABB validation.""" + from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams + from isaaclab_arena.relations.placement_validators import NoOverlapValidator + + source = DummyObject( + "source", + bounding_box=AxisAlignedBoundingBox(min_point=(-0.01, -0.01, -0.01), max_point=(0.01, 0.01, 0.01)), + ) + target = _make_box_obj("target", sx=0.05, sy=0.05, sz=0.05) + target.collision_mode = CollisionMode.MESH + positions = [ + {source: (0.0, 0.0, 0.0), target: (0.25, 0.0, 0.0)}, + {source: (0.0, 0.0, 0.0), target: (0.25, 0.0, 0.0)}, + ] + source_bboxes = [ + AxisAlignedBoundingBox(min_point=(-0.01, -0.01, -0.01), max_point=(0.01, 0.01, 0.01)), + AxisAlignedBoundingBox(min_point=(-0.3, -0.3, -0.01), max_point=(0.3, 0.3, 0.01)), + ] + bboxes = [{source: source_bbox, target: target.get_bounding_box()} for source_bbox in source_bboxes] + params = ObjectPlacerParams( + solver_params=RelationSolverParams(collision_mode=CollisionMode.BBOX, clearance_m=0.0, verbose=False) + ) + monkeypatch.setattr(NoOverlapValidator, "_spheres_penetrate_mesh", lambda *args, **kwargs: False) + + assert NoOverlapValidator(params).validate_batch(positions, [{}, {}], bboxes, []) == [True, False] + + +@requires_warp +def test_mesh_subject_with_meshless_anchor_uses_aabb_validation(): + """A meshless anchor forces AABB validation.""" + from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams + from isaaclab_arena.relations.placement_validators import NoOverlapValidator + + anchor = DummyObject( + "anchor", + bounding_box=AxisAlignedBoundingBox(min_point=(-0.1, -0.1, -0.1), max_point=(0.1, 0.1, 0.1)), + ) + anchor.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + anchor.add_relation(IsAnchor()) + robot = DummyObject( + "robot", + bounding_box=AxisAlignedBoundingBox(min_point=(-0.1, -0.1, -0.1), max_point=(0.1, 0.1, 0.1)), + collision_mesh=trimesh.creation.box(extents=(2.0, 2.0, 2.0)), + ) + positions = {anchor: (0.0, 0.0, 0.0), robot: (0.5, 0.0, 0.0)} + bboxes = {anchor: anchor.get_bounding_box(), robot: robot.get_bounding_box()} + params = ObjectPlacerParams( + solver_params=RelationSolverParams(collision_mode=CollisionMode.MESH, clearance_m=0.0, verbose=False) + ) + + assert NoOverlapValidator(params)._validate(positions, bboxes, orientations={}, collision_objects=[]) + + @requires_warp def test_yawed_aabb_proxy_validation_is_not_double_rotated(): """AABB proxy spheres built from yaw-expanded bboxes must not rotate by source yaw again.""" diff --git a/isaaclab_arena/tests/test_no_collision_loss.py b/isaaclab_arena/tests/test_no_collision_loss.py index 3b17b5fb33..ad50dafc7b 100644 --- a/isaaclab_arena/tests/test_no_collision_loss.py +++ b/isaaclab_arena/tests/test_no_collision_loss.py @@ -8,7 +8,6 @@ import math import torch -from isaaclab_arena.relations.loss_primitives import interval_overlap_axis_loss from isaaclab_arena.relations.relation_loss_strategies import NoCollisionLossStrategy from isaaclab_arena.relations.relation_solver import RelationSolver from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -94,27 +93,17 @@ def _single_pair_no_overlap_loss( parent_world_bbox: AxisAlignedBoundingBox, ) -> torch.Tensor: """Single-pair no-overlap loss; the reference the vectorized solver path must reproduce.""" - single_input = child_pos.dim() == 1 - if single_input: - child_pos = child_pos.unsqueeze(0) - - c = clearance_m - parent_x_min = parent_world_bbox.min_point[:, 0] - c - parent_x_max = parent_world_bbox.max_point[:, 0] + c - parent_y_min = parent_world_bbox.min_point[:, 1] - c - parent_y_max = parent_world_bbox.max_point[:, 1] + c - parent_z_min = parent_world_bbox.min_point[:, 2] - c - parent_z_max = parent_world_bbox.max_point[:, 2] + c - - child_world_min = child_pos + child_bbox.min_point - child_world_max = child_pos + child_bbox.max_point - - overlap_x = interval_overlap_axis_loss(child_world_min[:, 0], child_world_max[:, 0], parent_x_min, parent_x_max) - overlap_y = interval_overlap_axis_loss(child_world_min[:, 1], child_world_max[:, 1], parent_y_min, parent_y_max) - overlap_z = interval_overlap_axis_loss(child_world_min[:, 2], child_world_max[:, 2], parent_z_min, parent_z_max) - - total_loss = slope * (overlap_x * overlap_y * overlap_z) - return total_loss.squeeze(0) if single_input else total_loss + child_pos = child_pos.reshape(-1, 3) + subject_min = (child_pos + child_bbox.min_point).unsqueeze(0) + subject_max = (child_pos + child_bbox.max_point).unsqueeze(0) + strategy = NoCollisionLossStrategy(slope=slope) + return strategy.compute_loss_batched( + clearance_m, + subject_min, + subject_max, + parent_world_bbox.min_point.unsqueeze(0), + parent_world_bbox.max_point.unsqueeze(0), + ).squeeze() # ============================================================================= @@ -198,7 +187,7 @@ def test_no_collision_positive_loss_when_3d_overlap(): def test_no_collision_loss_scales_with_slope(): - """Test that NoCollision loss scales with slope (loss = slope * overlap_volume).""" + """Test that normalized penetration loss scales with slope.""" box_a = _create_box("box_a") box_b = _create_box("box_b") @@ -214,15 +203,15 @@ def test_no_collision_loss_scales_with_slope(): assert torch.isclose(loss_20, 2.0 * loss_10, rtol=1e-5) -def test_no_collision_loss_volume_formula(): - """Test that NoCollision loss equals slope * overlap volume for known overlap (clearance_m=0).""" +def test_no_collision_loss_normalized_depth_formula(): + """Loss is slope times shortest separating depth divided by subject size.""" box_a = _create_box("box_a", size=0.2) box_b = _create_box("box_b", size=0.2) child_pos = torch.tensor([0.1, 0.1, 0.1]) parent_world_bbox = box_b.get_bounding_box().translated((0.15, 0.15, 0.15)) - # Overlap [0.15, 0.3]^3, volume 0.15^3. Expected loss = 10 * 0.15^3. - expected_loss = 10.0 * (0.15**3) + # The nearest separating face is 0.15 m away; subject max extent is 0.2 m. + expected_loss = 10.0 * 0.15 / 0.2 loss = _single_pair_no_overlap_loss( 10.0, clearance_m=0.0, child_pos=child_pos, child_bbox=box_a.bounding_box, parent_world_bbox=parent_world_bbox @@ -230,6 +219,34 @@ def test_no_collision_loss_volume_formula(): assert torch.isclose(loss, torch.tensor(expected_loss), rtol=1e-4) +def test_no_collision_loss_is_scale_invariant(): + """Equal fractional penetration has equal loss at different scales.""" + strategy = NoCollisionLossStrategy(slope=10.0) + + def loss_for_size(size: float) -> torch.Tensor: + subject_min = torch.tensor([[[0.0, 0.0, 0.0]]]) + subject_max = torch.tensor([[[size, size, size]]]) + obstacle_min = torch.tensor([[[0.5 * size, 0.0, 0.0]]]) + obstacle_max = torch.tensor([[[1.5 * size, size, size]]]) + return strategy.compute_loss_batched(0.0, subject_min, subject_max, obstacle_min, obstacle_max) + + torch.testing.assert_close(loss_for_size(0.1), loss_for_size(2.0)) + + +def test_no_collision_containment_has_translation_gradient(): + """Containment retains a translation gradient.""" + strategy = NoCollisionLossStrategy(slope=10.0) + subject_min = torch.tensor([[[0.0, 0.0, 0.0]]], requires_grad=True) + subject_max = subject_min + 1.0 + obstacle_min = torch.tensor([[[-2.0, -2.0, -2.0]]]) + obstacle_max = torch.tensor([[[2.0, 2.0, 2.0]]]) + + strategy.compute_loss_batched(0.0, subject_min, subject_max, obstacle_min, obstacle_max).sum().backward() + + assert subject_min.grad is not None + assert torch.count_nonzero(subject_min.grad) > 0 + + # ============================================================================= # RelationSolver with built-in no-overlap tests # ============================================================================= @@ -626,7 +643,7 @@ def test_compute_loss_batched_direct(): loss = strategy.compute_loss_batched(0.0, subject_min, subject_max, obstacle_min, obstacle_max) assert loss.shape == (2, 1) - assert torch.isclose(loss[0, 0], torch.tensor(10.0 * 0.1**3), rtol=1e-4) # slope * overlap volume + assert torch.isclose(loss[0, 0], torch.tensor(5.0), rtol=1e-4) assert torch.isclose(loss[1, 0], torch.tensor(0.0), atol=1e-6) diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index 29b7dfbeda..025a010659 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -149,16 +149,11 @@ def scene_getitem(self, name: str) -> MagicMock: return env -def _solve_and_place_with_pool(env, env_ids, objects, pool): +def _solve_and_place_with_pool(env, env_ids, pool): """Call the reset event with the same runtime params EventTermCfg stores.""" from isaaclab_arena.relations.placement_events import solve_and_place_objects - return solve_and_place_objects( - env, - env_ids, - assets=objects, - placement_pool=pool, - ) + return solve_and_place_objects(env, env_ids, placement_pool=pool) def test_solve_and_place_objects_writes_poses_to_sim(): @@ -176,7 +171,7 @@ def test_solve_and_place_objects_writes_poses_to_sim(): placer_params = ObjectPlacerParams(solver_params=solver_params) pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=10) - _solve_and_place_with_pool(env, env_ids, objects, pool) + _solve_and_place_with_pool(env, env_ids, pool) # Anchor (desk) should NOT have been written. assert "desk" not in env._assets, "Anchor pose should not be written to sim" @@ -206,6 +201,7 @@ def test_solve_and_place_objects_uses_runtime_pool(): class Pool: num_envs = 1 + objects = [desk, robot] def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: assert env_ids == [0] @@ -221,7 +217,6 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: solve_and_place_objects( env, torch.tensor([0]), - assets=[desk, robot], placement_pool=Pool(), ) @@ -296,7 +291,7 @@ def test_reset_placement_asset_pose_per_env_requires_full_env_coverage(): def test_get_placement_pool_returns_runtime_pool(): - from isaaclab_arena.relations.placement_events import get_placement_pool + from isaaclab_arena.relations.placement_events import PlacementPoolHandle, get_placement_pool class Pool: pass @@ -306,6 +301,9 @@ class Pool: env.unwrapped.event_manager.get_term_cfg.return_value.params = {"placement_pool": pool} assert get_placement_pool(env) is pool + env.unwrapped.event_manager.get_term_cfg.return_value.params = {"placement_pool": PlacementPoolHandle(pool)} + assert get_placement_pool(env) is pool + def test_solve_and_place_objects_applies_random_yaw(): """With random_yaw_init enabled the runtime path should write yawed (non-identity) poses.""" @@ -328,7 +326,7 @@ def test_solve_and_place_objects_applies_random_yaw(): ) pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=10) - _solve_and_place_with_pool(env, env_ids, objects, pool) + _solve_and_place_with_pool(env, env_ids, pool) # Anchor (desk) is never rotated or written, even with random yaw enabled. assert "desk" not in env._assets, "Anchor pose should not be written to sim" @@ -356,7 +354,7 @@ def test_solve_and_place_objects_skips_empty_env_ids(): placer_params = ObjectPlacerParams(solver_params=solver_params) pool = PooledObjectPlacer(objects=[desk, box1, box2], placer_params=placer_params, pool_size=10) - _solve_and_place_with_pool(env, torch.tensor([], dtype=torch.int64), [desk, box1, box2], pool) + _solve_and_place_with_pool(env, torch.tensor([], dtype=torch.int64), pool) assert len(env._assets) == 0, "No writes should occur for empty env_ids" @@ -373,7 +371,7 @@ def test_solve_and_place_objects_skips_none_env_ids(): placer_params = ObjectPlacerParams(solver_params=solver_params) pool = PooledObjectPlacer(objects=[desk, box1, box2], placer_params=placer_params, pool_size=10) - _solve_and_place_with_pool(env, None, [desk, box1, box2], pool) + _solve_and_place_with_pool(env, None, pool) assert len(env._assets) == 0, "No writes should occur for None env_ids" @@ -394,7 +392,7 @@ def test_solve_and_place_objects_handles_multiple_env_ids(): placer_params = ObjectPlacerParams(solver_params=solver_params) pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=12, num_envs=num_envs) - _solve_and_place_with_pool(env, env_ids, objects, pool) + _solve_and_place_with_pool(env, env_ids, pool) assert "desk" not in env._assets, "Anchor pose should not be written to sim" @@ -424,7 +422,7 @@ def test_solve_and_place_objects_partial_reset_homogeneous_pool_consumes_only_re pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=12, num_envs=num_envs) available_before = pool.total_remaining - _solve_and_place_with_pool(env, env_ids, objects, pool) + _solve_and_place_with_pool(env, env_ids, pool) available_after = pool.total_remaining assert available_before - available_after == len(env_ids) @@ -436,11 +434,11 @@ def test_solve_and_place_objects_writes_invalid_fallback_layout(capsys): from isaaclab_arena.relations.placement_result import PlacementResult desk, box1, box2 = _create_test_objects() - objects = [desk, box1, box2] env = _make_mock_env(num_envs=1) class InvalidPool: num_envs = 1 + objects = [desk, box1, box2] def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: assert env_ids == [0] @@ -453,7 +451,7 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: ) } - _solve_and_place_with_pool(env, torch.tensor([0]), objects, InvalidPool()) + _solve_and_place_with_pool(env, torch.tensor([0]), InvalidPool()) captured = capsys.readouterr() assert set(env._assets) == {box1.name, box2.name} @@ -466,12 +464,12 @@ def test_solve_and_place_objects_partial_reset_applies_absolute_env_origin(): from isaaclab_arena.relations.placement_result import PlacementResult desk, box1, box2 = _create_test_objects() - objects = [desk, box1, box2] env = _make_mock_env(num_envs=4) env.scene.env_origins[2] = torch.tensor([10.0, 0.0, 0.0]) class EnvIndexedPool: num_envs = 4 + objects = [desk, box1, box2] requested_env_ids = None def sample_without_replacement(self, count: int) -> list[PlacementResult]: @@ -493,7 +491,7 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: } pool = EnvIndexedPool() - _solve_and_place_with_pool(env, torch.tensor([2]), objects, pool) + _solve_and_place_with_pool(env, torch.tensor([2]), pool) box1_pose = env._assets[box1.name].write_root_pose_to_sim.call_args[0][0] box2_pose = env._assets[box2.name].write_root_pose_to_sim.call_args[0][0] @@ -510,14 +508,14 @@ def test_solve_and_place_objects_asserts_env_indexed_pool_size_matches_scene(): """Env-indexed pool slots must line up with absolute Isaac Lab env ids.""" desk, box1, box2 = _create_test_objects() - objects = [desk, box1, box2] env = _make_mock_env(num_envs=2) class MismatchedEnvIndexedPool: num_envs = 1 + objects = [desk, box1, box2] with pytest.raises(AssertionError, match="scene has 2 env origins"): - _solve_and_place_with_pool(env, torch.tensor([0]), objects, MismatchedEnvIndexedPool()) + _solve_and_place_with_pool(env, torch.tensor([0]), MismatchedEnvIndexedPool()) def test_pooled_placer_sample_without_replacement_returns_different_layouts(): @@ -982,6 +980,13 @@ def test_solve_and_apply_relation_placement_drops_embodiment_from_event_params() assert params.reachability_config.embodiment is embodiment # ...while the pool the reset event captured no longer references the embodiment -- on the placer params # and on every built validator alike -- so configclass never deep-copies or recurses into it. - pool = event.params["placement_pool"] + from isaaclab.utils.configclass import _validate + + from isaaclab_arena.relations.placement_events import PlacementPoolHandle, resolve_placement_pool + + pool_handle = event.params["placement_pool"] + assert isinstance(pool_handle, PlacementPoolHandle) + _validate(event, prefix="") + pool = resolve_placement_pool(pool_handle) assert pool._placer.params.reachability_config.embodiment is None assert all(v._params.reachability_config.embodiment is None for v in pool._placer._validators) diff --git a/isaaclab_arena/tests/test_relation_solver_background_collision.py b/isaaclab_arena/tests/test_relation_solver_background_collision.py index 83dfb443ac..9ee8eda831 100644 --- a/isaaclab_arena/tests/test_relation_solver_background_collision.py +++ b/isaaclab_arena/tests/test_relation_solver_background_collision.py @@ -76,6 +76,23 @@ def _mesh_box(name: str, extents: tuple[float, float, float], position: tuple[fl return obj +def _make_usd_background(): + """Background stub for USD mesh extraction tests.""" + from isaaclab_arena.assets.background import Background + from isaaclab_arena.assets.object_base import ObjectType + from isaaclab_arena.utils.pose import Pose + + background = Background.__new__(Background) + background.name = "kitchen" + background.usd_path = "/tmp/kitchen.usda" + background.scale = (1.0, 1.0, 1.0) + background.object_type = ObjectType.BASE + background.collision_mode = None + background.initial_pose = Pose.identity() + background.repair_collision_mesh_non_watertight = False + return background + + def test_background_collision_object_combines_meshes(): """Multiple fixed background meshes are represented as one collision-only object.""" from isaaclab_arena.relations.background_collision_object import FixedCollisionObject, make_fixed_collision_objects @@ -125,7 +142,7 @@ def test_background_collision_objects_reject_failed_whole_background(monkeypatch monkeypatch.setattr( WarpMeshAndSphereCache, "get_collision_mesh", - lambda self, obj: left.get_collision_mesh() if obj is left else None, + lambda self, obj, excluded_prim_paths=(): left.get_collision_mesh() if obj is left else None, ) with pytest.raises(AssertionError, match="whole-scene Background"): @@ -147,7 +164,7 @@ def test_warp_mesh_cache_caches_unsupported_usd_geometry(monkeypatch): obj.repair_collision_mesh_non_watertight = True calls = {"count": 0} - def fail_extract(usd_path, scale): + def fail_extract(usd_path, scale, excluded_prim_paths=()): calls["count"] += 1 raise UnsupportedCollisionGeometryError("Unsupported non-mesh geometry in /tmp/kitchen.usd: /World/cube") @@ -159,6 +176,116 @@ def fail_extract(usd_path, scale): assert calls["count"] == 1 +def test_warp_mesh_cache_keys_exclusions(monkeypatch): + """Different anchor exclusions cannot reuse a stale whole-background mesh.""" + from isaaclab_arena.assets.object import Object + from isaaclab_arena.assets.object_base import ObjectType + from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache + + obj = Object.__new__(Object) + obj.name = "kitchen" + obj.usd_path = "/tmp/kitchen.usda" + obj.scale = (1.0, 1.0, 1.0) + obj.object_type = ObjectType.BASE + obj.repair_collision_mesh_non_watertight = False + calls = [] + + def fake_extract(usd_path, scale, excluded_prim_paths=()): + calls.append(tuple(excluded_prim_paths)) + return _mesh_box("mesh", (0.2, 0.2, 0.2), (0.0, 0.0, 0.0)).get_collision_mesh() + + monkeypatch.setattr("isaaclab_arena.utils.usd_helpers.extract_trimesh_from_usd", fake_extract) + manager = WarpMeshAndSphereCache(device="cpu") + + manager.get_collision_mesh(obj, excluded_prim_paths=["/Kitchen/counter"]) + manager.get_collision_mesh(obj, excluded_prim_paths=["/Kitchen/counter"]) + manager.get_collision_mesh(obj, excluded_prim_paths=["/Kitchen/floor"]) + + assert calls == [("/Kitchen/counter",), ("/Kitchen/floor",)] + + +def test_background_anchor_exclusions_can_remove_all_meshes(monkeypatch): + """A background with no mesh left after filtering contributes no obstacle.""" + from isaaclab_arena.relations.background_collision_object import make_fixed_collision_objects + from isaaclab_arena.utils.usd_helpers import AllCollisionMeshesExcludedError + + kitchen = _make_usd_background() + + def fake_extract(usd_path, scale, excluded_prim_paths=()): + assert excluded_prim_paths == ("/Kitchen",) + raise AllCollisionMeshesExcludedError("all geometry excluded") + + monkeypatch.setattr("isaaclab_arena.utils.usd_helpers.extract_trimesh_from_usd", fake_extract) + + assert ( + make_fixed_collision_objects( + [kitchen], + excluded_prim_paths_by_object={kitchen: ["/Kitchen"]}, + ) + == [] + ) + + +def test_background_exclusions_preserve_unsupported_geometry_error(monkeypatch): + """Unsupported geometry outside exclusions remains a fatal background extraction failure.""" + import pytest + + from isaaclab_arena.relations.background_collision_object import make_fixed_collision_objects + from isaaclab_arena.utils.usd_helpers import UnsupportedCollisionGeometryError + + kitchen = _make_usd_background() + + def fail_extract(usd_path, scale, excluded_prim_paths=()): + raise UnsupportedCollisionGeometryError("unsupported geometry remains") + + monkeypatch.setattr("isaaclab_arena.utils.usd_helpers.extract_trimesh_from_usd", fail_extract) + + with pytest.raises(AssertionError, match="whole-scene Background"): + make_fixed_collision_objects( + [kitchen], + excluded_prim_paths_by_object={kitchen: ["/Kitchen/counter"]}, + ) + + +def test_passive_background_excludes_relation_anchor_subtrees(monkeypatch): + """Background aggregation omits geometry represented separately by relation anchors.""" + from unittest.mock import MagicMock + + import isaaclab_arena.relations.passive_collision_objects as passive_module + from isaaclab_arena.assets.background import Background + from isaaclab_arena.assets.object_reference import ObjectReference + from isaaclab_arena.relations.passive_collision_objects import get_passive_collision_objects + from isaaclab_arena.utils.pose import Pose + + kitchen = Background.__new__(Background) + kitchen.name = "kitchen" + kitchen.usd_path = "/tmp/kitchen.usda" + kitchen.initial_pose = Pose.identity() + kitchen.relations = [] + + counter = MagicMock(spec=ObjectReference) + counter.parent_asset = kitchen + counter.get_prim_path_in_parent_usd.return_value = "/Kitchen/counter" + calls = {} + + def fake_make_fixed(objects, excluded_prim_paths_by_object=None): + calls["objects"] = list(objects) + calls["exclusions"] = excluded_prim_paths_by_object + return list(objects) + + monkeypatch.setattr(passive_module, "make_fixed_collision_objects", fake_make_fixed) + + result = get_passive_collision_objects( + [kitchen], + include_background=True, + background_mesh_exclusions=[counter], + ) + + assert result == [kitchen] + assert calls["objects"] == [kitchen] + assert calls["exclusions"] == {kitchen: ["/Kitchen/counter"]} + + def test_background_collision_objects_treat_background_none_pose_as_identity(monkeypatch): """A Background with no initial pose is fixed at the USD origin for mesh aggregation.""" import torch @@ -173,7 +300,9 @@ def test_background_collision_objects_treat_background_none_pose_as_identity(mon kitchen.initial_pose = None mesh_source = _mesh_box("source", (0.2, 0.2, 0.2), (0.0, 0.0, 0.0)) monkeypatch.setattr( - WarpMeshAndSphereCache, "get_collision_mesh", lambda self, obj: mesh_source.get_collision_mesh() + WarpMeshAndSphereCache, + "get_collision_mesh", + lambda self, obj, excluded_prim_paths=(): mesh_source.get_collision_mesh(), ) collision_objects = make_fixed_collision_objects([kitchen]) @@ -392,7 +521,9 @@ def fake_object(name, relations, usd_path, pose, spec=Object): } original = passive_collision_module.make_fixed_collision_objects - passive_collision_module.make_fixed_collision_objects = lambda objects: list(objects) + passive_collision_module.make_fixed_collision_objects = lambda objects, excluded_prim_paths_by_object=None: list( + objects + ) try: no_combine = passive_collision_module.get_passive_collision_objects(scene.assets.values()) combined = passive_collision_module.get_passive_collision_objects( @@ -644,27 +775,28 @@ def fake_solve_and_apply_relation_placement(*args, **kwargs): assert calls["objects"] == [embodiment] -def test_relation_placement_includes_background_mesh_for_object_mesh_override(monkeypatch): - """Object-level MESH override enables aggregate background meshes.""" +def test_relation_placement_forwards_anchor_background_mesh_exclusions(monkeypatch): + """Anchored object references are excluded from their background mesh.""" + from unittest.mock import MagicMock + import isaaclab_arena.environments.relation_solver_interface as interface_module + from isaaclab_arena.assets.object_reference import ObjectReference from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.relation_solver_params import CollisionMode, RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor - from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox - mesh_object = DummyObject( - "mesh_object", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.1, -0.1, -0.1), max_point=(0.1, 0.1, 0.1)), - ) - mesh_object.add_relation(IsAnchor()) - mesh_object.collision_mode = CollisionMode.MESH + reference = MagicMock(spec=ObjectReference) + reference.name = "counter" + reference.collision_mode = CollisionMode.MESH + reference.get_scene_key.return_value = "counter" + reference.get_relations.return_value = [IsAnchor()] calls = {} - def fake_get_passive_collision_objects(assets, include_background: bool = False): + def fake_get_passive_collision_objects(assets, include_background: bool = False, background_mesh_exclusions=()): calls["assets"] = list(assets) calls["include_background"] = include_background + calls["background_mesh_exclusions"] = list(background_mesh_exclusions) return [] class FakePooledObjectPlacer: @@ -678,11 +810,12 @@ def __init__(self, objects, placer_params, pool_size, num_envs, collision_object monkeypatch.setattr(interface_module, "PooledObjectPlacer", FakePooledObjectPlacer) placer_params = ObjectPlacerParams(solver_params=RelationSolverParams(collision_mode=CollisionMode.BBOX)) - solve_and_apply_relation_placement([mesh_object], num_envs=1, placer_params=placer_params, scene_assets=[]) + solve_and_apply_relation_placement([reference], num_envs=1, placer_params=placer_params, scene_assets=[]) assert calls["assets"] == [] assert calls["include_background"] is True - assert calls["objects"] == [mesh_object] + assert calls["background_mesh_exclusions"] == [reference] + assert calls["objects"] == [reference] assert calls["collision_objects"] == [] @@ -706,7 +839,7 @@ def test_relation_placement_includes_background_mesh_for_background_override(mon placed_object.add_relation(IsAnchor()) calls = {} - def fake_get_passive_collision_objects(assets, include_background: bool = False): + def fake_get_passive_collision_objects(assets, include_background: bool = False, background_mesh_exclusions=()): calls["assets"] = list(assets) calls["include_background"] = include_background return [] @@ -752,7 +885,7 @@ def test_relation_placement_skips_background_mesh_for_default_bbox(monkeypatch): placed_object.add_relation(IsAnchor()) calls = {} - def fake_get_passive_collision_objects(assets, include_background: bool = False): + def fake_get_passive_collision_objects(assets, include_background: bool = False, background_mesh_exclusions=()): calls["assets"] = list(assets) calls["include_background"] = include_background return [] diff --git a/isaaclab_arena/tests/test_relation_solver_interface.py b/isaaclab_arena/tests/test_relation_solver_interface.py index 561a446770..2be63964b7 100644 --- a/isaaclab_arena/tests/test_relation_solver_interface.py +++ b/isaaclab_arena/tests/test_relation_solver_interface.py @@ -34,8 +34,13 @@ def _make_box(name: str = "box"): class _FakePlacementPool: - def __init__(self, layouts) -> None: + def __init__(self, layouts, objects=None) -> None: self._layouts = layouts + self._objects = objects or [] + + @property + def objects(self): + return self._objects def sample_with_replacement(self, count: int): return self._layouts[:count] @@ -142,7 +147,10 @@ def test_dynamic_spawn_pose_event_params_use_runtime_assets(): desk = _make_desk() box = _make_box() - placement_pool = _FakePlacementPool([_fallback_layout(positions={box: (0.1, 0.2, 0.3)})]) + placement_pool = _FakePlacementPool( + [_fallback_layout(positions={box: (0.1, 0.2, 0.3)})], + objects=[desk, box], + ) event_cfg = _apply_dynamic_spawn_pose( assets=[desk, box], @@ -150,8 +158,48 @@ def test_dynamic_spawn_pose_event_params_use_runtime_assets(): anchor_assets={desk}, ) - assert [asset.name for asset in event_cfg.params["assets"]] == ["desk", "box"] assert "placement_pool" in event_cfg.params + from isaaclab_arena.relations.placement_events import PlacementPoolHandle, resolve_placement_pool + + assert isinstance(event_cfg.params["placement_pool"], PlacementPoolHandle) + assert [asset.name for asset in resolve_placement_pool(event_cfg.params["placement_pool"]).objects] == [ + "desk", + "box", + ] + + +def test_dynamic_spawn_pose_event_cfg_deepcopy_after_mesh_solve(): + """EventTermCfg deep-copies params; handle shares the pool after mesh caches are released.""" + import copy + import trimesh + + from isaaclab.managers import EventTermCfg + + from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement + from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams + from isaaclab_arena.relations.relation_solver_params import CollisionMode, RelationSolverParams + from isaaclab_arena.relations.relations import On + + desk = _make_desk() + box = _make_box() + box.add_relation(On(desk, clearance_m=0.01)) + box.collision_mode = CollisionMode.MESH + box._collision_mesh = trimesh.creation.box(extents=(0.2, 0.2, 0.2)) + + params = ObjectPlacerParams( + placement_seed=17, + resolve_on_reset=True, + min_unique_layouts_per_env=1, + solver_params=RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=50), + ) + event_cfg = solve_and_apply_relation_placement([desk, box], num_envs=1, placer_params=params) + + assert event_cfg is not None + assert isinstance(event_cfg, EventTermCfg) + copy.deepcopy(event_cfg) + from isaaclab.utils.configclass import _validate + + _validate(event_cfg, prefix="") def test_static_embodiment_placement_stores_per_env_poses(): diff --git a/isaaclab_arena/tests/test_usd_helpers.py b/isaaclab_arena/tests/test_usd_helpers.py index 4a7007f280..1336677e00 100644 --- a/isaaclab_arena/tests/test_usd_helpers.py +++ b/isaaclab_arena/tests/test_usd_helpers.py @@ -115,6 +115,38 @@ def _test_compute_local_bounding_box_from_usd_prim_path(simulation_app, asset_di return True +def _test_mesh_exclusion_errors(simulation_app, asset_dir: pathlib.Path) -> bool: + """Distinguish fully excluded meshes from malformed included meshes.""" + import pytest + from pxr import Gf, Usd, UsdGeom + + from isaaclab_arena.utils.usd_helpers import ( + AllCollisionMeshesExcludedError, + NoCollisionMeshError, + extract_trimesh_from_usd, + ) + + usd_path = asset_dir / "mesh_exclusions.usda" + stage = Usd.Stage.CreateNew(usd_path.as_posix()) + root = UsdGeom.Xform.Define(stage, "/Root") + stage.SetDefaultPrim(root.GetPrim()) + excluded_mesh = UsdGeom.Mesh.Define(stage, "/Root/Excluded") + excluded_mesh.GetPointsAttr().Set([Gf.Vec3f(0, 0, 0), Gf.Vec3f(1, 0, 0), Gf.Vec3f(0, 1, 0)]) + excluded_mesh.GetFaceVertexCountsAttr().Set([3]) + excluded_mesh.GetFaceVertexIndicesAttr().Set([0, 1, 2]) + stage.GetRootLayer().Save() + + with pytest.raises(AllCollisionMeshesExcludedError): + extract_trimesh_from_usd(usd_path.as_posix(), excluded_prim_paths=["/Root/Excluded"]) + + UsdGeom.Mesh.Define(stage, "/Root/Malformed") + stage.GetRootLayer().Save() + with pytest.raises(NoCollisionMeshError) as error: + extract_trimesh_from_usd(usd_path.as_posix(), excluded_prim_paths=["/Root/Excluded"]) + assert type(error.value) is NoCollisionMeshError + return True + + def test_compute_local_bounding_box_from_usd(tmp_path: pathlib.Path): result = run_simulation_app_function( _test_compute_local_bounding_box_from_usd, @@ -131,3 +163,11 @@ def test_compute_local_bounding_box_from_usd_prim_path(tmp_path: pathlib.Path): asset_dir=tmp_path, ) assert result, "Test failed" + + +def test_mesh_exclusion_errors(tmp_path: pathlib.Path): + assert run_simulation_app_function( + _test_mesh_exclusion_errors, + headless=HEADLESS, + asset_dir=tmp_path, + ) diff --git a/isaaclab_arena/utils/usd_helpers.py b/isaaclab_arena/utils/usd_helpers.py index 3e25076188..2cbd3f3d0d 100644 --- a/isaaclab_arena/utils/usd_helpers.py +++ b/isaaclab_arena/utils/usd_helpers.py @@ -7,6 +7,7 @@ import numpy as np import trimesh +from collections.abc import Sequence from contextlib import contextmanager from pxr import Gf, Usd, UsdGeom, UsdLux, UsdPhysics @@ -23,6 +24,10 @@ class UnsupportedCollisionGeometryError(NoCollisionMeshError): """USD geometry exists but cannot be represented as a collision mesh.""" +class AllCollisionMeshesExcludedError(NoCollisionMeshError): + """Every collision mesh is under an excluded USD subtree.""" + + def get_all_prims( stage: Usd.Stage, prim: Usd.Prim | None = None, prims_list: list[Usd.Prim] | None = None ) -> list[Usd.Prim]: @@ -282,6 +287,7 @@ def compute_local_bounding_box_from_prim( def extract_trimesh_from_usd( usd_path: str, scale: tuple[float, float, float] = (1.0, 1.0, 1.0), + excluded_prim_paths: Sequence[str] = (), ) -> trimesh.Trimesh: """Extract all UsdGeom.Mesh prims from a USD into a single trimesh. @@ -292,6 +298,7 @@ def extract_trimesh_from_usd( Args: usd_path: Path to the .usd/.usda/.usdc file. scale: (sx, sy, sz) per-axis scale factors applied in local frame. + excluded_prim_paths: Absolute USD prim paths whose complete subtrees are omitted. Returns: Combined trimesh with per-prim world transforms baked in. @@ -299,6 +306,10 @@ def extract_trimesh_from_usd( assert all( s > 0 for s in scale ), f"All scale components must be positive (negative scale flips winding/SDF sign), got {scale}" + excluded_paths = tuple(path.rstrip("/") for path in excluded_prim_paths) + assert all( + path.startswith("/") for path in excluded_paths + ), f"excluded_prim_paths must be absolute USD paths, got {excluded_paths}" stage = Usd.Stage.Open(usd_path) if stage is None: @@ -307,13 +318,20 @@ def extract_trimesh_from_usd( all_verts: list[np.ndarray] = [] all_faces: list[list[int]] = [] skipped_gprims: list[str] = [] + excluded_mesh_count = 0 + included_mesh_count = 0 offset = 0 for prim in stage.Traverse(): + prim_path = str(prim.GetPath()) + if any(prim_path == path or prim_path.startswith(f"{path}/") for path in excluded_paths): + excluded_mesh_count += int(prim.IsA(UsdGeom.Mesh)) + continue if not prim.IsA(UsdGeom.Mesh): if prim.IsA(UsdGeom.Gprim): skipped_gprims.append(str(prim.GetPath())) continue + included_mesh_count += 1 mesh_prim = UsdGeom.Mesh(prim) points = mesh_prim.GetPointsAttr().Get() face_vertex_counts = mesh_prim.GetFaceVertexCountsAttr().Get() @@ -351,6 +369,8 @@ def extract_trimesh_from_usd( raise UnsupportedCollisionGeometryError( f"Unsupported non-mesh geometry in {usd_path}: {', '.join(skipped_gprims)}" ) + if excluded_mesh_count and not included_mesh_count: + raise AllCollisionMeshesExcludedError(f"All mesh geometry excluded from {usd_path}") raise NoCollisionMeshError(f"No mesh geometry found in {usd_path}")