From 9e96027f3353c111496d3d3b426ec1f0900986e5 Mon Sep 17 00:00:00 2001 From: zhx06 Date: Mon, 27 Jul 2026 22:59:48 -0700 Subject: [PATCH] introduce oriented bbox Signed-off-by: zhx06 --- isaaclab_arena/assets/object.py | 6 +- isaaclab_arena/assets/object_reference.py | 50 +- isaaclab_arena/assets/object_set.py | 26 +- isaaclab_arena/embodiments/embodiment_base.py | 4 +- .../relations/background_collision_object.py | 6 +- .../relations/bounding_box_helpers.py | 83 +-- isaaclab_arena/relations/collision_mode.py | 2 +- isaaclab_arena/relations/collision_object.py | 6 +- isaaclab_arena/relations/loss_primitives.py | 38 -- isaaclab_arena/relations/mesh_pair_cache.py | 87 +-- isaaclab_arena/relations/no_overlap_mesh.py | 365 +++++----- .../{no_overlap_aabb.py => no_overlap_obb.py} | 92 +-- isaaclab_arena/relations/object_placer.py | 252 +++---- isaaclab_arena/relations/placement_asset.py | 17 +- isaaclab_arena/relations/placement_events.py | 6 +- isaaclab_arena/relations/placement_result.py | 4 +- .../relations/placement_validators.py | 369 +++++------ .../relations/relation_loss_strategies.py | 156 ++--- isaaclab_arena/relations/relation_solver.py | 85 ++- .../relations/relation_solver_params.py | 4 + .../relations/relation_solver_state.py | 141 ++-- isaaclab_arena/relations/warp_mesh_manager.py | 19 - isaaclab_arena/relations/warp_sdf_kernels.py | 10 - isaaclab_arena/tests/dummy_embodiment.py | 6 +- isaaclab_arena/tests/dummy_object.py | 8 +- isaaclab_arena/tests/test_bounding_box.py | 478 ++++++++----- .../tests/test_embodiment_collision_mesh.py | 2 +- isaaclab_arena/tests/test_face_to.py | 102 ++- .../tests/test_heterogeneous_placement.py | 111 ++-- .../tests/test_isaac_sim_debug_draw.py | 27 + isaaclab_arena/tests/test_mesh_collision.py | 324 +++++---- .../tests/test_no_collision_loss.py | 147 ++-- .../tests/test_object_mass_variation.py | 6 +- .../tests/test_object_placer_init.py | 74 ++- .../test_object_placer_reproducibility.py | 47 +- isaaclab_arena/tests/test_object_set.py | 14 +- isaaclab_arena/tests/test_placement_events.py | 42 +- isaaclab_arena/tests/test_pose.py | 10 + isaaclab_arena/tests/test_position_limits.py | 8 +- .../tests/test_reference_objects.py | 213 +++++- .../tests/test_relation_loss_strategies.py | 14 +- ...st_relation_solver_background_collision.py | 47 +- .../tests/test_relation_solver_embodiment.py | 13 +- .../tests/test_relation_solver_interface.py | 22 +- isaaclab_arena/tests/test_usd_helpers.py | 4 +- isaaclab_arena/tests/test_usd_pose_helpers.py | 24 + .../tests/test_usd_scale_helpers.py | 29 +- .../tests/test_validate_placement.py | 73 +- isaaclab_arena/utils/bounding_box.py | 627 ++++++++---------- isaaclab_arena/utils/isaac_sim_debug_draw.py | 58 +- isaaclab_arena/utils/trimesh.py | 6 +- isaaclab_arena/utils/usd_helpers.py | 150 +---- isaaclab_arena/utils/usd_pose_helpers.py | 6 +- isaaclab_arena/utils/yaw.py | 6 +- .../ik_reachability_validator.py | 54 +- isaaclab_arena_curobo/ik_solver.py | 4 +- .../tests/test_ik_reachability_validator.py | 131 +++- .../utils/ik_solver_utils.py | 30 +- .../review_gui/simapp/asset_usd.py | 3 +- .../relations/dummy_object_placer_notebook.py | 40 +- .../relations/example_object.py | 6 +- .../relation_solver_visualization_notebook.py | 30 +- .../relations/relation_solver_visualizer.py | 60 +- 63 files changed, 2580 insertions(+), 2304 deletions(-) rename isaaclab_arena/relations/{no_overlap_aabb.py => no_overlap_obb.py} (69%) diff --git a/isaaclab_arena/assets/object.py b/isaaclab_arena/assets/object.py index 5ccab0f9ca..ff5133888b 100644 --- a/isaaclab_arena/assets/object.py +++ b/isaaclab_arena/assets/object.py @@ -15,7 +15,7 @@ from isaaclab_arena.assets.object_base import ObjectBase, ObjectType from isaaclab_arena.assets.object_utils import detect_object_type from isaaclab_arena.relations.relations import RelationBase -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose from isaaclab_arena.utils.usd.rigid_bodies import find_shallowest_rigid_body from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd, has_light, open_stage @@ -63,7 +63,7 @@ def __init__( self.object_cfg = self._init_object_cfg() self._pose_event_cfg = self._build_reset_event() - def get_bounding_box(self) -> AxisAlignedBoundingBox: + def get_bounding_box(self) -> OrientedBoundingBox: """Get local bounding box (relative to object origin).""" assert self.usd_path is not None if self.bounding_box is None: @@ -74,7 +74,7 @@ def get_corners(self, pos: torch.Tensor) -> torch.Tensor: assert self.usd_path is not None if self.bounding_box is None: self.bounding_box = compute_local_bounding_box_from_usd(self.usd_path, self.scale) - return self.bounding_box.get_corners_at(pos) + return self.bounding_box.translated(pos).get_corners() def is_initial_pose_set(self) -> bool: return self.initial_pose is not None diff --git a/isaaclab_arena/assets/object_reference.py b/isaaclab_arena/assets/object_reference.py index 9310b07c4d..6c100b09b6 100644 --- a/isaaclab_arena/assets/object_reference.py +++ b/isaaclab_arena/assets/object_reference.py @@ -3,17 +3,19 @@ # # SPDX-License-Identifier: Apache-2.0 +import torch import trimesh from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg from isaaclab.sensors.contact_sensor.contact_sensor_cfg import ContactSensorCfg +from isaaclab.utils.math import matrix_from_quat from pxr import Usd from isaaclab_arena.affordances.openable import Openable from isaaclab_arena.assets.object import Object from isaaclab_arena.assets.object_base import ObjectBase, ObjectType from isaaclab_arena.relations.relations import IsAnchor, RelationBase -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, quaternion_to_90_deg_z_quarters +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose from isaaclab_arena.utils.usd_helpers import ( NoCollisionMeshError, @@ -28,13 +30,17 @@ class ObjectReference(ObjectBase): """An object which *refers* to an existing element in the scene""" def __init__(self, parent_asset: Object, **kwargs): + parent_scale = parent_asset.scale + assert all( + component > 0 for component in parent_scale + ), f"ObjectReference parent scale must be positive, got {parent_scale}." super().__init__(**kwargs) self.parent_asset = parent_asset - self._parent_scale = parent_asset.scale + self._parent_scale = parent_scale # Get the prim's transform pose (not geometry center - solver is origin-agnostic) self.initial_pose_relative_to_parent = self._get_referenced_prim_pose_relative_to_parent(parent_asset) self.object_cfg = self._init_object_cfg() - self._bounding_box: AxisAlignedBoundingBox | None = None + self._bounding_box: OrientedBoundingBox | None = None self._collision_mesh: trimesh.Trimesh | None = None # None is a valid cached result for meshless prims; this flag distinguishes that from not-yet-loaded. self._collision_mesh_loaded = False @@ -63,7 +69,7 @@ def add_relation(self, relation: RelationBase) -> None: ) self.relations.append(relation) - def get_bounding_box(self) -> AxisAlignedBoundingBox: + def get_bounding_box(self) -> OrientedBoundingBox: """Get local bounding box of the referenced prim (relative to prim transform). The bounding box is relative to the prim's transform origin, consistent with @@ -77,23 +83,18 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox: self.prim_path, self.parent_asset, parent_stage ) raw_bbox = compute_local_bounding_box_from_prim(parent_stage, prim_path_in_usd) - # Apply parent's scale (no centering - solver is origin-agnostic) - self._bounding_box = raw_bbox.scaled(self._parent_scale) + scaled_corners = self._transform_raw_local_points(raw_bbox.get_corners()) + self._bounding_box = OrientedBoundingBox.from_min_max( + min_point=scaled_corners.amin(dim=1), + max_point=scaled_corners.amax(dim=1), + ) return self._bounding_box - def get_world_bounding_box(self) -> AxisAlignedBoundingBox: - """Bounding box in world coordinates. - - get_bounding_box() is already axis-aligned in the parent's frame, so only the parent's - placement rotation (identity or a 90° Z multiple) and the prim's world position are applied. - """ + def get_world_bounding_box(self) -> OrientedBoundingBox: + """Return the referenced prim's bounding box in world coordinates.""" box = self.get_bounding_box() - world_position = self.get_initial_pose().position_xyz - parent_pose = self.parent_asset.initial_pose - if parent_pose is None: - return box.translated(world_position) - quarters = quaternion_to_90_deg_z_quarters(parent_pose.rotation_xyzw) - return box.rotated_90_around_z(quarters).translated(world_position) + world_pose = self.get_initial_pose() + return box.transformed(world_pose.position_xyz, world_pose.rotation_xyzw) def get_collision_mesh(self) -> trimesh.Trimesh | None: """Return the referenced prim's collision mesh in its local frame, or None if unavailable.""" @@ -118,7 +119,18 @@ def _extract_collision_mesh(self) -> trimesh.Trimesh: ) if not parent_stage.GetPrimAtPath(prim_path_in_usd): raise ValueError(f"No prim found with path {prim_path_in_usd} in {self.parent_asset.usd_path}") - return extract_trimesh_from_prim(parent_stage, prim_path_in_usd, self._parent_scale) + mesh = extract_trimesh_from_prim(parent_stage, prim_path_in_usd, (1.0, 1.0, 1.0)) + vertices = torch.as_tensor(mesh.vertices, dtype=torch.float64) + mesh.vertices = self._transform_raw_local_points(vertices).numpy() + return mesh + + def _transform_raw_local_points(self, points: torch.Tensor) -> torch.Tensor: + """Apply parent scale to referenced-local points while preserving the reference pose.""" + rotation = matrix_from_quat(points.new_tensor(self.initial_pose_relative_to_parent.rotation_xyzw).unsqueeze(0))[ + 0 + ] + scale = points.new_tensor(self._parent_scale) + return ((points @ rotation.T) * scale) @ rotation def get_contact_sensor_cfg(self, contact_against_object: ObjectBase | None = None) -> ContactSensorCfg: # NOTE(alexmillane): Right now this requires that the object diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py index 085cfc0b3c..a04699f7a2 100644 --- a/isaaclab_arena/assets/object_set.py +++ b/isaaclab_arena/assets/object_set.py @@ -12,7 +12,7 @@ from isaaclab_arena.assets.object import Object from isaaclab_arena.assets.object_base import ObjectBase, ObjectType from isaaclab_arena.assets.object_utils import detect_object_type -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose from isaaclab_arena.utils.usd.object_set_utils import rescale_rename_rigid_body_and_save_to_cache from isaaclab_arena.utils.usd.rigid_bodies import find_shallowest_rigid_body @@ -97,15 +97,15 @@ def object_usd_paths(self) -> list[str]: return [self.member_usd_paths[idx] for idx in self.variant_indices_by_env] return self.member_usd_paths - def get_bounding_box(self) -> AxisAlignedBoundingBox: + def get_bounding_box(self) -> OrientedBoundingBox: """Return one local bbox for callers that cannot vary by env. - The returned bbox has shape (1, 3) and uses the member with the - greatest z-extent. Heterogeneous placement uses + The returned bbox has N=1 and uses the member with the greatest + z-extent. Heterogeneous placement uses get_bounding_box_per_env() after assign_variants() so each env uses its actual variant geometry. """ - return max(self.objects, key=lambda obj: obj.get_bounding_box().size[0, 2].item()).get_bounding_box() + return max(self.objects, key=lambda obj: obj.get_bounding_box().half_extents[0, 2].item()).get_bounding_box() def assign_variants(self, num_envs: int, variant_seed: int | None = None) -> None: """Fix one member-variant index per environment. @@ -132,18 +132,17 @@ def assign_variants(self, num_envs: int, variant_seed: int | None = None) -> Non print(f"Warning: RigidObjectSet '{self.name}' regenerating variant assignments for {num_envs} envs.") self._set_variant_indices_by_env(self._generate_variant_indices(num_envs, variant_seed=variant_seed)) - def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox: + def get_bounding_box_per_env(self, num_envs: int) -> OrientedBoundingBox: """Return each env's actual variant bbox. - Requires assign_variants(num_envs) to have been called first. The - returned bbox has shape (num_envs, 3). + Requires assign_variants(num_envs) to have been called first. Args: num_envs: Number of environments. Must match the assignment. Returns: - AxisAlignedBoundingBox with min_point / max_point of - shape (num_envs, 3). + OrientedBoundingBox with center and half_extents shape + (num_envs, 3), and rotation_xyzw shape (num_envs, 4). """ assert self.variant_indices_by_env is not None, ( f"RigidObjectSet '{self.name}' has no variant assignment; " @@ -155,9 +154,10 @@ def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox: ) bounding_boxes = [obj.get_bounding_box() for obj in self.objects] - min_pts = torch.stack([bounding_boxes[idx].min_point[0] for idx in self.variant_indices_by_env], dim=0) - max_pts = torch.stack([bounding_boxes[idx].max_point[0] for idx in self.variant_indices_by_env], dim=0) - return AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts) + centers = torch.stack([bounding_boxes[idx].center[0] for idx in self.variant_indices_by_env], dim=0) + half_extents = torch.stack([bounding_boxes[idx].half_extents[0] for idx in self.variant_indices_by_env], dim=0) + rotations = torch.stack([bounding_boxes[idx].rotation_xyzw[0] for idx in self.variant_indices_by_env], dim=0) + return OrientedBoundingBox(center=centers, half_extents=half_extents, rotation_xyzw=rotations) def get_contact_sensor_cfg(self, contact_against_object: ObjectBase | None = None) -> ContactSensorCfg: # We assume that by here, our USDs have been modified to be compatible with each other diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index 9308a4268f..fde2317749 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -14,7 +14,7 @@ from isaaclab_arena.embodiments.common.arm_mode import ArmMode from isaaclab_arena.relations.placement_asset import PlaceableAsset -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.cameras import ArenaCameraCfg, make_camera_observation_cfg from isaaclab_arena.utils.configclass import combine_configclass_instances from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange @@ -59,7 +59,7 @@ def __init__( self._collision_mesh: trimesh.Trimesh | None = None """Lazily-extracted robot collision mesh, cached so the USD is opened once.""" - def get_bounding_box(self) -> AxisAlignedBoundingBox: + def get_bounding_box(self) -> OrientedBoundingBox: """Return root-relative bounds computed from the articulation's USD geometry.""" # Import locally because USD/pxr is available only after simulation initialization. from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd diff --git a/isaaclab_arena/relations/background_collision_object.py b/isaaclab_arena/relations/background_collision_object.py index d2590766a6..cbeccc8e74 100644 --- a/isaaclab_arena/relations/background_collision_object.py +++ b/isaaclab_arena/relations/background_collision_object.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING from isaaclab_arena.relations.collision_mode import CollisionMode -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose from isaaclab_arena.utils.trimesh import bounding_box_from_mesh, mesh_in_world_frame @@ -54,11 +54,11 @@ def get_initial_pose(self) -> Pose: """Return identity pose because the mesh is already baked into world coordinates.""" return self._pose - def get_bounding_box(self) -> AxisAlignedBoundingBox: + def get_bounding_box(self) -> OrientedBoundingBox: """Return the mesh bounds; identical to the world bounds since the mesh is in world frame.""" return self._bounding_box - def get_world_bounding_box(self) -> AxisAlignedBoundingBox: + def get_world_bounding_box(self) -> OrientedBoundingBox: """Return the mesh bounds in world frame.""" return self._bounding_box diff --git a/isaaclab_arena/relations/bounding_box_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py index 5d3b0ea9f7..9c9116651f 100644 --- a/isaaclab_arena/relations/bounding_box_helpers.py +++ b/isaaclab_arena/relations/bounding_box_helpers.py @@ -3,17 +3,14 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Bounding-box helpers for heterogeneous placement. - -Keeps num_envs and per-env geometry logic out of placement assets. -""" +"""Per-environment bounding-box helpers.""" from __future__ import annotations from dataclasses import dataclass from typing import TYPE_CHECKING -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox if TYPE_CHECKING: from isaaclab_arena.relations.placement_asset import PlaceableAsset @@ -44,88 +41,72 @@ def assign_variants_for_envs(objects: list[PlaceableAsset], num_envs: int, place variant_set_idx += 1 -def get_bounding_box_per_env(obj: PlaceableAsset, num_envs: int) -> AxisAlignedBoundingBox: - """Return bounding boxes expanded to (num_envs, 3). - - RigidObjectSet delegates to its own get_bounding_box_per_env. - All other objects broadcast their single bbox. - """ +def get_bounding_box_per_env(obj: PlaceableAsset, num_envs: int) -> OrientedBoundingBox: + """Return one local bounding box per environment.""" from isaaclab_arena.assets.object_set import RigidObjectSet if isinstance(obj, RigidObjectSet): return obj.get_bounding_box_per_env(num_envs) bbox = obj.get_bounding_box() - return AxisAlignedBoundingBox( - min_point=bbox.min_point.expand(num_envs, 3), - max_point=bbox.max_point.expand(num_envs, 3), + return OrientedBoundingBox( + center=bbox.center.expand(num_envs, 3), + half_extents=bbox.half_extents.expand(num_envs, 3), + rotation_xyzw=bbox.rotation_xyzw.expand(num_envs, 4), ) @dataclass(frozen=True) class PerEnvBoundingBoxes: - """Per-env object bboxes, exposed in three layouts: + """Local object bounding boxes for each environment.""" - - get_bounding_boxes_for_env_id: one dict for a single env, bboxes (1, 3). - - get_bounding_boxes_for_all_envs: list[dict] of length num_envs, each bbox (1, 3). - - get_bounding_boxes_for_solver_candidates: one dict tiled to - (num_envs * candidates_per_env, 3), grouped contiguously by env. - """ + object_bboxes: dict[PlaceableAsset, OrientedBoundingBox] + """Boxes with center/half-extents shape (N, 3) and rotation shape (N, 4).""" - object_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox] num_envs: int + """Number of environments N.""" def __post_init__(self) -> None: assert self.num_envs >= 1, f"num_envs must be >= 1, got {self.num_envs}" for obj, bbox in self.object_bboxes.items(): assert ( - bbox.min_point.shape[0] == self.num_envs - ), f"Object '{obj.name}' bbox min_point has {bbox.min_point.shape[0]} envs, expected {self.num_envs}." + bbox.center.shape[0] == self.num_envs + ), f"Object '{obj.name}' bbox center has {bbox.center.shape[0]} envs, expected {self.num_envs}." assert ( - bbox.max_point.shape[0] == self.num_envs - ), f"Object '{obj.name}' bbox max_point has {bbox.max_point.shape[0]} envs, expected {self.num_envs}." - - def get_bounding_boxes_for_env_id(self, env_id: int) -> dict[PlaceableAsset, AxisAlignedBoundingBox]: - """Return object bboxes for a single env (each (1, 3)), used for per-env initialization and validation.""" - return { - obj: AxisAlignedBoundingBox( - min_point=bbox.min_point[env_id : env_id + 1], - max_point=bbox.max_point[env_id : env_id + 1], + bbox.half_extents.shape[0] == self.num_envs + ), f"Object '{obj.name}' bbox half_extents has {bbox.half_extents.shape[0]} envs, expected {self.num_envs}." + assert bbox.rotation_xyzw.shape[0] == self.num_envs, ( + f"Object '{obj.name}' bbox rotation_xyzw has {bbox.rotation_xyzw.shape[0]} envs, expected" + f" {self.num_envs}." ) - for obj, bbox in self.object_bboxes.items() - } - def get_bounding_boxes_for_all_envs(self) -> list[dict[PlaceableAsset, AxisAlignedBoundingBox]]: - """Return one-env bbox dicts for every env. + def get_bounding_boxes_for_env_id(self, env_id: int) -> dict[PlaceableAsset, OrientedBoundingBox]: + """Return object bboxes for one env with N=1.""" + return {obj: bbox[env_id] for obj, bbox in self.object_bboxes.items()} - The outer list has length num_envs. Each bbox has min_point/max_point - shape (1, 3). - """ + def get_bounding_boxes_for_all_envs(self) -> list[dict[PlaceableAsset, OrientedBoundingBox]]: + """Return num_envs one-env bbox dicts, each with N=1.""" return [self.get_bounding_boxes_for_env_id(env_id) for env_id in range(self.num_envs)] def get_bounding_boxes_for_solver_candidates( self, candidates_per_env: int - ) -> dict[PlaceableAsset, AxisAlignedBoundingBox]: + ) -> dict[PlaceableAsset, OrientedBoundingBox]: """Return bboxes tiled to one row per solver candidate. - Each bbox has shape (num_envs * candidates_per_env, 3). Rows are grouped - contiguously by env: rows [i * candidates_per_env : (i + 1) * candidates_per_env] - all hold env i's bbox. Callers recover the env via candidate_idx // candidates_per_env. + Each bbox has N=num_envs * candidates_per_env. Rows are grouped contiguously + by env; callers recover the env via candidate_idx // candidates_per_env. """ return { - obj: AxisAlignedBoundingBox( - min_point=bbox.min_point.repeat_interleave(candidates_per_env, dim=0), - max_point=bbox.max_point.repeat_interleave(candidates_per_env, dim=0), + obj: OrientedBoundingBox( + center=bbox.center.repeat_interleave(candidates_per_env, dim=0), + half_extents=bbox.half_extents.repeat_interleave(candidates_per_env, dim=0), + rotation_xyzw=bbox.rotation_xyzw.repeat_interleave(candidates_per_env, dim=0), ) for obj, bbox in self.object_bboxes.items() } def build_per_env_bounding_boxes(objects: list[PlaceableAsset], num_envs: int) -> PerEnvBoundingBoxes: - """Build per-env base bboxes for each placement object. - - Object orientation (marker roll/pitch/yaw plus sampled and FaceTo yaw) is applied later per - candidate in ObjectPlacer._rotate_candidate_bboxes, so these boxes carry object geometry only. - """ + """Build per-env local OBB geometry for each placement object.""" object_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects} return PerEnvBoundingBoxes(object_bboxes=object_bboxes, num_envs=num_envs) diff --git a/isaaclab_arena/relations/collision_mode.py b/isaaclab_arena/relations/collision_mode.py index d3efc9e72d..a285fb317f 100644 --- a/isaaclab_arena/relations/collision_mode.py +++ b/isaaclab_arena/relations/collision_mode.py @@ -16,7 +16,7 @@ class CollisionMode(Enum): """Collision-detection method for no-overlap constraints.""" BBOX = "bbox" - """Axis-aligned bounding box overlap volume (fast, conservative).""" + """Oriented bounding box penetration (fast).""" MESH = "mesh" """Sphere-to-SDF queries against actual mesh geometry (accurate, slower).""" diff --git a/isaaclab_arena/relations/collision_object.py b/isaaclab_arena/relations/collision_object.py index 115670f0a1..3e03d48cde 100644 --- a/isaaclab_arena/relations/collision_object.py +++ b/isaaclab_arena/relations/collision_object.py @@ -11,7 +11,7 @@ from typing import Protocol from isaaclab_arena.relations.collision_mode import CollisionMode -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange @@ -29,10 +29,10 @@ def is_anchor(self) -> bool: def get_initial_pose(self) -> Pose | PoseRange | PosePerEnv | None: raise NotImplementedError - def get_bounding_box(self) -> AxisAlignedBoundingBox: + def get_bounding_box(self) -> OrientedBoundingBox: raise NotImplementedError - def get_world_bounding_box(self) -> AxisAlignedBoundingBox: + def get_world_bounding_box(self) -> OrientedBoundingBox: raise NotImplementedError def get_collision_mesh(self) -> trimesh.Trimesh | None: diff --git a/isaaclab_arena/relations/loss_primitives.py b/isaaclab_arena/relations/loss_primitives.py index cc136564ee..fa70954191 100644 --- a/isaaclab_arena/relations/loss_primitives.py +++ b/isaaclab_arena/relations/loss_primitives.py @@ -122,41 +122,3 @@ def single_point_linear_loss( target = torch.tensor(target, dtype=value.dtype, device=value.device) return slope * torch.abs(value - target) - - -def interval_overlap_axis_loss( - child_min: torch.Tensor, - child_max: torch.Tensor, - parent_min: torch.Tensor | float, - parent_max: torch.Tensor | float, - slope: float = 1.0, -) -> torch.Tensor: - """ReLU-style interval overlap: zero when separated, slope * overlap length otherwise. - - Used by NoCollisionLossStrategy for per-axis overlap. Intervals [child_min, child_max] - and [parent_min, parent_max]; loss is zero when they do not overlap, else - slope * overlap_length. - - Args: - child_min: Child interval min (tensor for gradient flow). - child_max: Child interval max (tensor). - parent_min: Parent interval min (tensor or float). - parent_max: Parent interval max (tensor or float). - slope: Gradient magnitude (default: 1.0). - - Returns: - Zero when intervals are separated; otherwise slope * overlap length. - """ - assert isinstance(child_min, torch.Tensor), f"child_min must be a torch.Tensor, got {type(child_min)}" - assert isinstance(child_max, torch.Tensor), f"child_max must be a torch.Tensor, got {type(child_max)}" - assert torch.all(child_min <= child_max), "child_min must be <= child_max for valid interval [child_min, child_max]" - if not isinstance(parent_min, torch.Tensor): - parent_min = torch.tensor(parent_min, dtype=child_min.dtype, device=child_min.device) - if not isinstance(parent_max, torch.Tensor): - parent_max = torch.tensor(parent_max, dtype=child_max.dtype, device=child_max.device) - assert torch.all( - parent_min <= parent_max - ), "parent_min must be <= parent_max for valid interval [parent_min, parent_max]" - overlap_high = torch.minimum(child_max, parent_max) - overlap_low = torch.maximum(child_min, parent_min) - return single_boundary_linear_loss(overlap_high, overlap_low, slope=slope, penalty_side="greater") diff --git a/isaaclab_arena/relations/mesh_pair_cache.py b/isaaclab_arena/relations/mesh_pair_cache.py index 3598e75234..cda55a9bc2 100644 --- a/isaaclab_arena/relations/mesh_pair_cache.py +++ b/isaaclab_arena/relations/mesh_pair_cache.py @@ -19,10 +19,7 @@ class MeshPairEntry(NamedTuple): - """One directed sphere-to-mesh collision pair (subject spheres vs obstacle mesh). - - Dimensions: S = sphere count for this pair's subject, B = batch_size. - """ + """One directed sphere-to-mesh collision pair.""" subject: PlaceableAsset """Subject (sphere source) object.""" @@ -36,35 +33,14 @@ class MeshPairEntry(NamedTuple): fixed_obstacle_pos: torch.Tensor | None """(3,) world-frame position of the fixed obstacle; None for non-fixed obstacles.""" - fixed_obstacle_yaw: float - """Fixed obstacle Z-yaw in radians (0.0 for non-fixed obstacles).""" + fixed_obstacle_rotation: torch.Tensor | None + """(4,) world-frame xyzw rotation of the fixed obstacle; None for non-fixed obstacles.""" centers_local: torch.Tensor """(S, 3) sphere centers in subject-local frame.""" - subject_applies_yaw: bool - """True when subject sphere centers are in the object's unrotated local frame.""" - radii: torch.Tensor - """(S,) sphere radii.""" - - subject_bbox_min: torch.Tensor - """(B, 3) subject bounding box min corners.""" - - subject_bbox_max: torch.Tensor - """(B, 3) subject bounding box max corners.""" - - subject_bbox_includes_yaw: bool - """True when subject bbox extents are already yaw-expanded.""" - - obstacle_bbox_min: torch.Tensor - """(B, 3) obstacle bounding box min corners.""" - - obstacle_bbox_max: torch.Tensor - """(B, 3) obstacle bounding box max corners.""" - - obstacle_bbox_includes_yaw: bool - """True when obstacle bbox extents are already yaw-expanded.""" + """(S,) asset-local sphere radii.""" warp_mesh: wp.Mesh """Warp mesh asset for the obstacle.""" @@ -72,13 +48,7 @@ class MeshPairEntry(NamedTuple): @dataclass(slots=True) class MeshPairCache: - """Precomputed per-pair collision data for the vectorized multi-mesh kernel. - - Dimensions: P = num_pairs (ordered subject/obstacle pairs), B = batch_size (num envs), - S = total_spheres (sum of sphere counts across all P pairs; each subject object is decomposed - into multiple covering spheres via greedy_sphere_decomposition), - M = num_unique_meshes (distinct collision meshes referenced by the pairs). - """ + """Precomputed data for P directed pairs, S spheres, and M target meshes.""" all_centers_local: torch.Tensor """(S, 3) sphere centers in each subject's local frame, concatenated across pairs.""" @@ -92,35 +62,14 @@ class MeshPairCache: pair_obstacle_objs: list[PlaceableAsset | CollisionObject] """(P,) obstacle (mesh target) object reference per pair.""" - pair_subject_applies_yaw: list[bool] - """(P,) True when subject sphere centers are in the subject's unrotated local frame.""" - pair_obstacle_is_fixed: list[bool] """(P,) True if the obstacle is fixed in world coordinates.""" pair_fixed_obstacle_pos: list[torch.Tensor | None] """(P,) world position for fixed obstacles (None for non-fixed obstacles).""" - pair_fixed_obstacle_yaw: list[float] - """(P,) fixed obstacle yaw in radians (0.0 for non-fixed obstacles).""" - - pair_subject_bbox_min: torch.Tensor - """(P, B, 3) subject bounding box min corners.""" - - pair_subject_bbox_max: torch.Tensor - """(P, B, 3) subject bounding box max corners.""" - - pair_subject_bbox_includes_yaw: list[bool] - """(P,) True when subject bbox extents are already yaw-expanded.""" - - pair_obstacle_bbox_min: torch.Tensor - """(P, B, 3) obstacle bounding box min corners.""" - - pair_obstacle_bbox_max: torch.Tensor - """(P, B, 3) obstacle bounding box max corners.""" - - pair_obstacle_bbox_includes_yaw: list[bool] - """(P,) True when obstacle bbox extents are already yaw-expanded.""" + pair_fixed_obstacle_rotation: list[torch.Tensor | None] + """(P,) fixed-obstacle world xyzw rotations, each (4,); None for dynamic obstacles.""" pair_max_radius: torch.Tensor """(P,) maximum sphere radius across all spheres in each pair.""" @@ -146,18 +95,22 @@ class MeshPairCache: def __post_init__(self) -> None: assert len(self.pair_subject_objs) == self.num_pairs, "pair_subject_objs length mismatch" assert len(self.pair_obstacle_objs) == self.num_pairs, "pair_obstacle_objs length mismatch" - assert len(self.pair_subject_applies_yaw) == self.num_pairs, "pair_subject_applies_yaw length mismatch" - assert ( - len(self.pair_subject_bbox_includes_yaw) == self.num_pairs - ), "pair_subject_bbox_includes_yaw length mismatch" - assert ( - len(self.pair_obstacle_bbox_includes_yaw) == self.num_pairs - ), "pair_obstacle_bbox_includes_yaw length mismatch" assert len(self.pair_obstacle_is_fixed) == self.num_pairs, "pair_obstacle_is_fixed length mismatch" + assert len(self.pair_fixed_obstacle_pos) == self.num_pairs, "pair_fixed_obstacle_pos length mismatch" + assert len(self.pair_fixed_obstacle_rotation) == self.num_pairs, "pair_fixed_obstacle_rotation length mismatch" assert self.all_centers_local.shape[0] == self.total_spheres, "all_centers_local size mismatch" assert self.all_radii.shape[0] == self.total_spheres, "all_radii size mismatch" assert self.sphere_pair_id.shape[0] == self.total_spheres, "sphere_pair_id size mismatch" assert self.sphere_mesh_idx.shape[0] == self.total_spheres, "sphere_mesh_idx size mismatch" assert int(self.pair_sphere_count.sum().item()) == self.total_spheres, "pair_sphere_count sum mismatch" - for i, (is_fixed, pos) in enumerate(zip(self.pair_obstacle_is_fixed, self.pair_fixed_obstacle_pos)): - assert not is_fixed or pos is not None, f"pair {i}: obstacle_is_fixed=True but fixed_obstacle_pos is None" + for i, (is_fixed, pos, rotation) in enumerate( + zip( + self.pair_obstacle_is_fixed, + self.pair_fixed_obstacle_pos, + self.pair_fixed_obstacle_rotation, + strict=True, + ) + ): + assert not is_fixed or ( + pos is not None and rotation is not None + ), f"pair {i}: fixed obstacles require position and rotation" diff --git a/isaaclab_arena/relations/no_overlap_mesh.py b/isaaclab_arena/relations/no_overlap_mesh.py index 056e6f23f5..dbfcbd6651 100644 --- a/isaaclab_arena/relations/no_overlap_mesh.py +++ b/isaaclab_arena/relations/no_overlap_mesh.py @@ -13,26 +13,64 @@ from typing import TYPE_CHECKING import warp as wp +from isaaclab.utils.math import quat_apply, quat_apply_inverse from isaaclab_arena.relations.collision_mode import CollisionMode, object_uses_mesh_collision from isaaclab_arena.relations.mesh_pair_cache import MeshPairCache, MeshPairEntry +from isaaclab_arena.relations.placement_asset import PlaceableAsset from isaaclab_arena.relations.relation_solver_state import RelationSolverState -from isaaclab_arena.relations.warp_sdf_kernels import clamp_sdf_sentinel, multi_mesh_sdf +from isaaclab_arena.relations.warp_sdf_kernels import has_sdf_sentinel, multi_mesh_sdf +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose -from isaaclab_arena.utils.yaw import rotate_points_by_yaw_batch, yaw_from_quat_xyzw if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject - 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 + + +def transform_points_between_frames( + points_source: torch.Tensor, + source_position: torch.Tensor, + source_rotation_xyzw: torch.Tensor, + target_position: torch.Tensor, + target_rotation_xyzw: torch.Tensor, +) -> torch.Tensor: + """Transform source-local points into a target-local frame using xyzw quaternions.""" + count = points_source.shape[0] + source_rotation = source_rotation_xyzw.reshape(-1, 4) + target_rotation = target_rotation_xyzw.reshape(-1, 4) + if source_rotation.shape[0] == 1: + source_rotation = source_rotation.expand(count, 4) + if target_rotation.shape[0] == 1: + target_rotation = target_rotation.expand(count, 4) + assert source_rotation.shape[0] == count and target_rotation.shape[0] == count + points_world = quat_apply(source_rotation, points_source) + source_position + return quat_apply_inverse(target_rotation, points_world - target_position) + + +def _bbox_for_batch(bbox: OrientedBoundingBox, batch_size: int) -> OrientedBoundingBox: + """Broadcast one asset-local box or return its complete candidate batch.""" + assert bbox.num_envs in (1, batch_size), f"Expected one box or B={batch_size}, got {bbox.num_envs}." + return OrientedBoundingBox.from_tensors_unchecked( + bbox.center.expand(batch_size, 3), + bbox.half_extents.expand(batch_size, 3), + bbox.rotation_xyzw.expand(batch_size, 4), + ) + + +def _concatenate_box_batches(boxes: list[OrientedBoundingBox]) -> OrientedBoundingBox: + """Concatenate pair-major box batches without revalidating tensors.""" + return OrientedBoundingBox.from_tensors_unchecked( + torch.cat([bbox.center for bbox in boxes], dim=0), + torch.cat([bbox.half_extents for bbox in boxes], dim=0), + torch.cat([bbox.rotation_xyzw for bbox in boxes], dim=0), + ) def compute_no_overlap_loss_mesh( state: RelationSolverState, mesh_cache: MeshPairCache | None, mesh_manager: WarpMeshAndSphereCache, - orientations: list[dict[PlaceableAsset, float]] | None, clearance_m: float, slope: float, debug: bool, @@ -42,8 +80,7 @@ def compute_no_overlap_loss_mesh( Args: state: Current solver state with positions and batch info. mesh_cache: Precomputed collision pair data (None = no pairs). - mesh_manager: Warp mesh/sphere cache (for sentinel warnings). - orientations: Per-env yaw angles per object. + mesh_manager: Warp mesh/sphere cache associated with the pair cache. clearance_m: Minimum clearance between objects. slope: Gradient magnitude for overlap loss. debug: Print per-pair loss when True. @@ -55,111 +92,83 @@ def compute_no_overlap_loss_mesh( return total_loss num_pairs = mesh_cache.num_pairs + batch_size = state.batch_size + position_cache = {obj: state.get_position(obj) for obj in state.optimizable_objects} + rotation_cache = {obj: state.get_rotation(obj) for obj in state.optimizable_objects} + + subject_positions = torch.stack([position_cache[obj] for obj in mesh_cache.pair_subject_objs]) + subject_rotations = torch.stack([rotation_cache[obj] for obj in mesh_cache.pair_subject_objs]) + obstacle_position_batches: list[torch.Tensor] = [] + obstacle_rotation_batches: list[torch.Tensor] = [] + obstacle_box_batches: list[OrientedBoundingBox] = [] + for p, obstacle in enumerate(mesh_cache.pair_obstacle_objs): + if mesh_cache.pair_obstacle_is_fixed[p]: + fixed_position = mesh_cache.pair_fixed_obstacle_pos[p] + fixed_rotation = mesh_cache.pair_fixed_obstacle_rotation[p] + assert fixed_position is not None and fixed_rotation is not None + obstacle_position_batches.append(fixed_position.expand(batch_size, 3)) + obstacle_rotation_batches.append(fixed_rotation.expand(batch_size, 4)) + obstacle_box_batches.append(_bbox_for_batch(state.get_fixed_obstacle_world_bbox(obstacle), batch_size)) + continue - # Per-env loop: per-env yaw and active-pair masking produce a different sphere subset per env. - for b in range(state.batch_size): - subject_positions = torch.stack( - [state.get_position(mesh_cache.pair_subject_objs[p])[b] for p in range(num_pairs)] + assert isinstance(obstacle, PlaceableAsset) + position = position_cache[obstacle].detach() + rotation = rotation_cache[obstacle].detach() + obstacle_position_batches.append(position) + obstacle_rotation_batches.append(rotation) + obstacle_box_batches.append( + _bbox_for_batch(state.get_base_bbox(obstacle), batch_size).transformed_unchecked(position, rotation) ) - obstacle_positions = torch.stack([ - ( - mesh_cache.pair_fixed_obstacle_pos[p] - if mesh_cache.pair_obstacle_is_fixed[p] - else state.get_position(mesh_cache.pair_obstacle_objs[p])[b].detach() - ) - for p in range(num_pairs) - ]) - - fixed_obstacle_yaws = mesh_cache.pair_fixed_obstacle_yaw - has_any_yaw = orientations is not None or any(y != 0.0 for y in fixed_obstacle_yaws) - if has_any_yaw: - ori_b = orientations[b] if orientations is not None else {} - subject_yaws = torch.tensor( - [ - ori_b.get(mesh_cache.pair_subject_objs[p], 0.0) if mesh_cache.pair_subject_applies_yaw[p] else 0.0 - for p in range(num_pairs) - ], - dtype=torch.float32, - device=device, - ) - obstacle_yaws = torch.tensor( - [ori_b.get(mesh_cache.pair_obstacle_objs[p], fixed_obstacle_yaws[p]) for p in range(num_pairs)], - dtype=torch.float32, - device=device, - ) - - # AABB overlap filter (yaw-aware): skip separated pairs. - margins = mesh_cache.pair_max_radius + clearance_m - s_bbox_min = mesh_cache.pair_subject_bbox_min[:, b, :] - s_bbox_max = mesh_cache.pair_subject_bbox_max[:, b, :] - o_bbox_min = mesh_cache.pair_obstacle_bbox_min[:, b, :] - o_bbox_max = mesh_cache.pair_obstacle_bbox_max[:, b, :] - - if has_any_yaw: - subject_bbox_yaws = torch.tensor( - [ - 0.0 if mesh_cache.pair_subject_bbox_includes_yaw[p] else subject_yaws[p].item() - for p in range(num_pairs) - ], - dtype=torch.float32, - device=device, - ) - obstacle_bbox_yaws = torch.tensor( - [ - 0.0 if mesh_cache.pair_obstacle_bbox_includes_yaw[p] else obstacle_yaws[p].item() - for p in range(num_pairs) - ], - dtype=torch.float32, - device=device, - ) - s_bbox_min, s_bbox_max = _rotate_bbox_extents(s_bbox_min, s_bbox_max, subject_bbox_yaws) - o_bbox_min, o_bbox_max = _rotate_bbox_extents(o_bbox_min, o_bbox_max, obstacle_bbox_yaws) - - subject_min = subject_positions + s_bbox_min - subject_max = subject_positions + s_bbox_max - obstacle_min = obstacle_positions + o_bbox_min - obstacle_max = obstacle_positions + o_bbox_max - - sep_subject = (subject_min - margins.unsqueeze(1)) > obstacle_max - sep_obstacle = (obstacle_min - margins.unsqueeze(1)) > subject_max - separated = sep_subject.any(dim=1) | sep_obstacle.any(dim=1) - active_pair = ~separated + obstacle_positions = torch.stack(obstacle_position_batches) + obstacle_rotations = torch.stack(obstacle_rotation_batches) + + subject_boxes = _concatenate_box_batches( + [_bbox_for_batch(state.get_base_bbox(obj), batch_size) for obj in mesh_cache.pair_subject_objs] + ).transformed_unchecked( + subject_positions.reshape(-1, 3), + subject_rotations.reshape(-1, 4), + ) + obstacle_boxes = _concatenate_box_batches(obstacle_box_batches) + subject_min, subject_max = subject_boxes.get_axis_aligned_bounds() + obstacle_min, obstacle_max = obstacle_boxes.get_axis_aligned_bounds() + broadphase_margin = (mesh_cache.pair_max_radius.repeat_interleave(batch_size) + clearance_m).unsqueeze(1) + separated = ((subject_min - broadphase_margin) > obstacle_max).any(dim=1) | ( + (obstacle_min - broadphase_margin) > subject_max + ).any(dim=1) + active_pair_env = ~separated + if not active_pair_env.any(): + return total_loss - if not active_pair.any(): - continue + environment_ids = torch.arange(batch_size, device=device, dtype=torch.long) + sphere_pair_env_ids = (mesh_cache.sphere_pair_id.unsqueeze(1) * batch_size + environment_ids.unsqueeze(0)).reshape( + -1 + ) + active_sphere_mask = active_pair_env[sphere_pair_env_ids] + active_pair_env_ids = sphere_pair_env_ids[active_sphere_mask] + active_centers = transform_points_between_frames( + mesh_cache.all_centers_local.repeat_interleave(batch_size, dim=0)[active_sphere_mask], + subject_positions.reshape(-1, 3)[active_pair_env_ids], + subject_rotations.reshape(-1, 4)[active_pair_env_ids], + obstacle_positions.reshape(-1, 3)[active_pair_env_ids], + obstacle_rotations.reshape(-1, 4)[active_pair_env_ids], + ) + active_radii = mesh_cache.all_radii.repeat_interleave(batch_size)[active_sphere_mask] + active_mesh_idx = mesh_cache.sphere_mesh_idx.repeat_interleave(batch_size)[active_sphere_mask].contiguous() + sdf_values = multi_mesh_sdf( + active_centers, + mesh_cache.mesh_id_array, + wp.from_torch(active_mesh_idx, dtype=wp.int32), + ) + assert not has_sdf_sentinel( + sdf_values + ), "MESH collision query could not resolve a target face; the promised collision geometry is unsupported." + penetration = torch.relu(active_radii + clearance_m - sdf_values) - offsets = subject_positions - obstacle_positions - sphere_active_mask = active_pair[mesh_cache.sphere_pair_id] - active_idx = sphere_active_mask.nonzero(as_tuple=True)[0] - - active_sphere_pair_id = mesh_cache.sphere_pair_id[active_idx] - local_centers = mesh_cache.all_centers_local[active_idx] - - # R(subject_yaw - obstacle_yaw) · local + R(-obstacle_yaw) · offset - if has_any_yaw: - net_yaws = (subject_yaws - obstacle_yaws)[active_sphere_pair_id] - local_centers = rotate_points_by_yaw_batch(local_centers, net_yaws) - - pair_offsets = offsets[active_sphere_pair_id] - obs_yaws = obstacle_yaws[active_sphere_pair_id] - rotated_offsets = rotate_points_by_yaw_batch(pair_offsets, -obs_yaws) - active_centers = local_centers + rotated_offsets - else: - active_centers = local_centers + offsets[active_sphere_pair_id] - active_radii = mesh_cache.all_radii[active_idx] - active_mesh_idx = mesh_cache.sphere_mesh_idx[active_idx].contiguous() - - active_mesh_indices_wp = wp.from_torch(active_mesh_idx, dtype=wp.int32) - sdf_values = multi_mesh_sdf(active_centers, mesh_cache.mesh_id_array, active_mesh_indices_wp) - 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 - active_pair_idx = active_pair.nonzero(as_tuple=True)[0] - total_loss[b] = total_loss[b] + slope * pair_mean[active_pair_idx].sum() + pair_env_sum = torch.zeros(num_pairs * batch_size, device=device, dtype=penetration.dtype) + pair_env_sum.index_add_(0, active_pair_env_ids, penetration) + pair_env_count = mesh_cache.pair_sphere_count.repeat_interleave(batch_size) + pair_env_mean = (pair_env_sum / pair_env_count).reshape(num_pairs, batch_size) + total_loss = slope * pair_env_mean.sum(dim=0) if debug: print(f" [NoOverlap MESH] total_loss={total_loss.tolist()}") @@ -173,7 +182,6 @@ def prepare_mesh_collision_cache( on_pairs: set[tuple[int, int]], warned_no_mesh: set[str], default_collision_mode: CollisionMode = CollisionMode.MESH, - bboxes_include_yaw: bool = False, ) -> MeshPairCache | None: """Precompute static per-pair mesh collision data. @@ -183,7 +191,6 @@ def prepare_mesh_collision_cache( on_pairs: Set of (id(a), id(b)) pairs linked by On relations (skipped). warned_no_mesh: Mutable set tracking which objects have already been warned about. default_collision_mode: Collision mode used by objects without a per-object override. - bboxes_include_yaw: True when state bboxes are already yaw-expanded. Returns: Combined MeshPairCache for all directed pairs, or None if no pairs qualify. @@ -202,7 +209,6 @@ def prepare_mesh_collision_cache( device, warned_no_mesh, default_collision_mode, - bboxes_include_yaw, ) return _finalize_mesh_cache(all_pairs, device) @@ -210,13 +216,12 @@ def prepare_mesh_collision_cache( def _collect_mesh_pairs( state: RelationSolverState, manager: WarpMeshAndSphereCache, - non_anchor_objects: list, + non_anchor_objects: list[PlaceableAsset], fixed_obstacles: list[PlaceableAsset | CollisionObject], on_pairs: set[tuple[int, int]], device: torch.device, warned_no_mesh: set[str], default_collision_mode: CollisionMode, - bboxes_include_yaw: bool, ) -> list[MeshPairEntry]: """Collect all directed mesh pairs (forward + reverse).""" pairs: list[MeshPairEntry] = [] @@ -224,22 +229,18 @@ def _collect_mesh_pairs( for i, child in enumerate(non_anchor_objects): child_uses_mesh = object_uses_mesh_collision(child, default_collision_mode) child_mesh = manager.get_collision_mesh(child) if child_uses_mesh else None - child_bbox = state.get_bbox(child).to(device) + child_bbox = state.get_base_bbox(child).to(device) child_bbox_is_invariant = child_bbox.is_batch_invariant() if child_uses_mesh and child_mesh is None and child.name not in warned_no_mesh: warned_no_mesh.add(child.name) fallback = ( - "using an AABB-sphere approximation for mesh-obstacle pairs" + "using an OBB-sphere approximation for mesh-obstacle pairs" if child_bbox_is_invariant - else "pair will use AABB fallback for varying per-env bboxes" + else "pair will use OBB fallback for varying per-env bboxes" ) print(f"[NoCollision] '{child.name}' has no collision mesh; {fallback}.") child_spheres = _get_subject_spheres(child_mesh, child_bbox, child, manager, device) - child_applies_yaw = child_mesh is not None or not bboxes_include_yaw - c_bbox_min = child_bbox.min_point.expand(state.batch_size, 3) - c_bbox_max = child_bbox.max_point.expand(state.batch_size, 3) - # child's spheres → fixed obstacle mesh (anchors plus passive background) for obstacle in fixed_obstacles: if (id(child), id(obstacle)) in on_pairs: continue @@ -251,107 +252,81 @@ def _collect_mesh_pairs( if obstacle_mesh is None: if object_uses_mesh_collision(obstacle, default_collision_mode) and obstacle.name not in warned_no_mesh: warned_no_mesh.add(obstacle.name) - print(f"[NoCollision] '{obstacle.name}' has no collision mesh; pair will use AABB fallback.") + print(f"[NoCollision] '{obstacle.name}' has no collision mesh; pair will use OBB fallback.") continue pose = obstacle.get_initial_pose() - assert pose is not None and isinstance( + assert isinstance( pose, Pose ), f"MESH collision requires fixed obstacle '{obstacle.name}' to have a fixed Pose initial_pose" - assert abs(pose.rotation_xyzw[0]) < 1e-6 and abs(pose.rotation_xyzw[1]) < 1e-6, ( - f"MESH collision requires fixed obstacle '{obstacle.name}' to have identity or " - f"pure-Z rotation, got rotation_xyzw={pose.rotation_xyzw}. " - "Roll/pitch fixed obstacles are not supported in MESH mode." - ) if child_spheres is None: continue - obstacle_bbox = obstacle.get_bounding_box().to(device) pairs.append( MeshPairEntry( subject=child, obstacle=obstacle, obstacle_is_fixed=True, fixed_obstacle_pos=torch.tensor(pose.position_xyz, dtype=torch.float32, device=device), - fixed_obstacle_yaw=yaw_from_quat_xyzw(pose.rotation_xyzw), + fixed_obstacle_rotation=torch.tensor(pose.rotation_xyzw, dtype=torch.float32, device=device), centers_local=child_spheres[:, :3], - subject_applies_yaw=child_applies_yaw, radii=child_spheres[:, 3], - subject_bbox_min=c_bbox_min, - subject_bbox_max=c_bbox_max, - subject_bbox_includes_yaw=bboxes_include_yaw, - obstacle_bbox_min=obstacle_bbox.min_point.expand(state.batch_size, 3), - obstacle_bbox_max=obstacle_bbox.max_point.expand(state.batch_size, 3), - obstacle_bbox_includes_yaw=False, warp_mesh=manager.get_warp_mesh(obstacle_mesh, obj=obstacle), ) ) - # Non-anchor pairs (bidirectional): forward + reverse for j in range(i + 1, len(non_anchor_objects)): other = non_anchor_objects[j] if (id(child), id(other)) in on_pairs: continue other_uses_mesh = object_uses_mesh_collision(other, default_collision_mode) other_mesh = manager.get_collision_mesh(other) if other_uses_mesh else None - other_bbox = state.get_bbox(other).to(device) + other_bbox = state.get_base_bbox(other).to(device) other_bbox_is_invariant = other_bbox.is_batch_invariant() if other_mesh is None and child_mesh is None: if other_uses_mesh and other.name not in warned_no_mesh: warned_no_mesh.add(other.name) fallback = ( - "using an AABB-sphere approximation for mesh-obstacle pairs" + "using an OBB-sphere approximation for mesh-obstacle pairs" if other_bbox_is_invariant - else "pair will use AABB fallback for varying per-env bboxes" + else "pair will use OBB fallback for varying per-env bboxes" ) print(f"[NoCollision] '{other.name}' has no collision mesh; {fallback}.") continue - o_bbox_min = other_bbox.min_point.expand(state.batch_size, 3) - o_bbox_max = other_bbox.max_point.expand(state.batch_size, 3) - - if other_mesh is not None and child_spheres is not None: - # forward: child's mesh/spheres or AABB-sphere approximation → other's mesh + child_target_mesh = child_mesh if child_mesh is not None else _bbox_proxy_mesh(child_bbox) + other_target_mesh = other_mesh if other_mesh is not None else _bbox_proxy_mesh(other_bbox) + if other_target_mesh is not None and child_spheres is not None: pairs.append( MeshPairEntry( subject=child, obstacle=other, obstacle_is_fixed=False, fixed_obstacle_pos=None, - fixed_obstacle_yaw=0.0, + fixed_obstacle_rotation=None, centers_local=child_spheres[:, :3], - subject_applies_yaw=child_applies_yaw, radii=child_spheres[:, 3], - subject_bbox_min=c_bbox_min, - subject_bbox_max=c_bbox_max, - subject_bbox_includes_yaw=bboxes_include_yaw, - obstacle_bbox_min=o_bbox_min, - obstacle_bbox_max=o_bbox_max, - obstacle_bbox_includes_yaw=bboxes_include_yaw, - warp_mesh=manager.get_warp_mesh(other_mesh, obj=other), + warp_mesh=manager.get_warp_mesh( + other_target_mesh, + obj=other if other_mesh is not None else None, + ), ) ) - if child_mesh is not None: - # reverse: other's mesh/spheres or AABB-sphere approximation → child's mesh + if child_target_mesh is not None: other_spheres = _get_subject_spheres(other_mesh, other_bbox, other, manager, device) if other_spheres is None: continue - other_applies_yaw = other_mesh is not None or not bboxes_include_yaw pairs.append( MeshPairEntry( subject=other, obstacle=child, obstacle_is_fixed=False, fixed_obstacle_pos=None, - fixed_obstacle_yaw=0.0, + fixed_obstacle_rotation=None, centers_local=other_spheres[:, :3], - subject_applies_yaw=other_applies_yaw, radii=other_spheres[:, 3], - subject_bbox_min=o_bbox_min, - subject_bbox_max=o_bbox_max, - subject_bbox_includes_yaw=bboxes_include_yaw, - obstacle_bbox_min=c_bbox_min, - obstacle_bbox_max=c_bbox_max, - obstacle_bbox_includes_yaw=bboxes_include_yaw, - warp_mesh=manager.get_warp_mesh(child_mesh, obj=child), + warp_mesh=manager.get_warp_mesh( + child_target_mesh, + obj=child if child_mesh is not None else None, + ), ) ) @@ -360,7 +335,7 @@ def _collect_mesh_pairs( def _get_subject_spheres( mesh: trimesh.Trimesh | None, - bbox: AxisAlignedBoundingBox, + bbox: OrientedBoundingBox, obj: PlaceableAsset, manager: WarpMeshAndSphereCache, device: torch.device, @@ -368,13 +343,24 @@ def _get_subject_spheres( """Return (S, 4) query spheres; return None for varying meshless bboxes.""" if mesh is not None: return manager.get_query_spheres(mesh, obj=obj).to(device) + box_mesh = _bbox_proxy_mesh(bbox) + if box_mesh is None: + return None + return manager.get_query_spheres(box_mesh).to(device) + + +def _bbox_proxy_mesh(bbox: OrientedBoundingBox) -> trimesh.Trimesh | None: + """Return one oriented box mesh for a batch-invariant OBB.""" if not bbox.is_batch_invariant(): return None center = bbox.center[0].detach().cpu().numpy() - extents = bbox.size[0].detach().cpu().numpy() + extents = (2.0 * bbox.half_extents[0]).detach().cpu().numpy() box_mesh = trimesh.creation.box(extents=extents) - box_mesh.apply_translation(center) - return manager.get_query_spheres(box_mesh).to(device) + rotation_xyzw = bbox.rotation_xyzw[0].detach().cpu().numpy() + transform = trimesh.transformations.quaternion_matrix(np.roll(rotation_xyzw, 1)) + transform[:3, 3] = center + box_mesh.apply_transform(transform) + return box_mesh def _finalize_mesh_cache(entries: list[MeshPairEntry], device: torch.device) -> MeshPairCache | None: @@ -385,9 +371,6 @@ def _finalize_mesh_cache(entries: list[MeshPairEntry], device: torch.device) -> mesh_id_map: dict[int, int] = {} mesh_id_values: list[int] = [] mesh_idx_per_sphere: list[int] = [] - pair_slices: list[tuple[int, int]] = [] - offset = 0 - for entry in entries: n_spheres = entry.centers_local.shape[0] mesh_key = id(entry.warp_mesh) @@ -395,49 +378,25 @@ def _finalize_mesh_cache(entries: list[MeshPairEntry], device: torch.device) -> mesh_id_map[mesh_key] = len(mesh_id_values) mesh_id_values.append(entry.warp_mesh.id) mesh_idx_per_sphere.extend([mesh_id_map[mesh_key]] * n_spheres) - pair_slices.append((offset, offset + n_spheres)) - offset += n_spheres - pair_sphere_count = torch.tensor([e - s for s, e in pair_slices], dtype=torch.float32, device=device) - sphere_pair_id = torch.repeat_interleave(torch.arange(len(pair_slices), device=device), pair_sphere_count.long()) + pair_sphere_count = torch.tensor( + [entry.centers_local.shape[0] for entry in entries], dtype=torch.float32, device=device + ) + sphere_pair_id = torch.repeat_interleave(torch.arange(len(entries), device=device), pair_sphere_count.long()) return MeshPairCache( all_centers_local=torch.cat([e.centers_local for e in entries], dim=0), all_radii=torch.cat([e.radii for e in entries], dim=0), pair_subject_objs=[e.subject for e in entries], pair_obstacle_objs=[e.obstacle for e in entries], - pair_subject_applies_yaw=[e.subject_applies_yaw for e in entries], pair_obstacle_is_fixed=[e.obstacle_is_fixed for e in entries], pair_fixed_obstacle_pos=[e.fixed_obstacle_pos for e in entries], - pair_fixed_obstacle_yaw=[e.fixed_obstacle_yaw for e in entries], - pair_subject_bbox_min=torch.stack([e.subject_bbox_min for e in entries]), - pair_subject_bbox_max=torch.stack([e.subject_bbox_max for e in entries]), - pair_subject_bbox_includes_yaw=[e.subject_bbox_includes_yaw for e in entries], - pair_obstacle_bbox_min=torch.stack([e.obstacle_bbox_min for e in entries]), - pair_obstacle_bbox_max=torch.stack([e.obstacle_bbox_max for e in entries]), - pair_obstacle_bbox_includes_yaw=[e.obstacle_bbox_includes_yaw for e in entries], + pair_fixed_obstacle_rotation=[e.fixed_obstacle_rotation for e in entries], pair_max_radius=torch.tensor([e.radii.max().item() for e in entries], device=device), sphere_pair_id=sphere_pair_id, sphere_mesh_idx=torch.tensor(mesh_idx_per_sphere, dtype=torch.int32, device=device), pair_sphere_count=pair_sphere_count, mesh_id_array=wp.array(np.array(mesh_id_values, dtype=np.uint64), dtype=wp.uint64, device=str(device)), num_pairs=len(entries), - total_spheres=offset, + total_spheres=int(pair_sphere_count.sum().item()), ) - - -def _rotate_bbox_extents( - bbox_min: torch.Tensor, bbox_max: torch.Tensor, yaws: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - """Return the AABB enclosing a Z-rotated bbox around the object origin.""" - min_x, min_y = bbox_min[:, 0], bbox_min[:, 1] - max_x, max_y = bbox_max[:, 0], bbox_max[:, 1] - corners_x = torch.stack([min_x, max_x, max_x, min_x], dim=1) - corners_y = torch.stack([min_y, min_y, max_y, max_y], dim=1) - cos_y = torch.cos(yaws).unsqueeze(1) - sin_y = torch.sin(yaws).unsqueeze(1) - rot_x = corners_x * cos_y - corners_y * sin_y - rot_y = corners_x * sin_y + corners_y * cos_y - rotated_min = torch.stack([rot_x.min(dim=1).values, rot_y.min(dim=1).values, bbox_min[:, 2]], dim=1) - rotated_max = torch.stack([rot_x.max(dim=1).values, rot_y.max(dim=1).values, bbox_max[:, 2]], dim=1) - return rotated_min, rotated_max diff --git a/isaaclab_arena/relations/no_overlap_aabb.py b/isaaclab_arena/relations/no_overlap_obb.py similarity index 69% rename from isaaclab_arena/relations/no_overlap_aabb.py rename to isaaclab_arena/relations/no_overlap_obb.py index a838c93457..b6d36ba79e 100644 --- a/isaaclab_arena/relations/no_overlap_aabb.py +++ b/isaaclab_arena/relations/no_overlap_obb.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""AABB-based no-overlap collision loss computation.""" +"""OBB-based no-overlap collision loss computation.""" from __future__ import annotations @@ -15,6 +15,7 @@ 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 +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject @@ -24,25 +25,19 @@ @dataclass(frozen=True) class NoOverlapPair: - """One directed overlap penalty: the subject box is pushed off the (detached) obstacle box. + """One directed overlap pair.""" - Dimensions: B = batch_size (num envs). - """ - - subject_min: torch.Tensor - """(B, 3) world-space min corner of the subject box.""" - - subject_max: torch.Tensor - """(B, 3) world-space max corner of the subject box.""" + subject: OrientedBoundingBox + """Subject world boxes; center and half-extents have shape (N, 3).""" - obstacle_min: torch.Tensor - """(B, 3) world-space min corner of the obstacle box.""" + obstacle: OrientedBoundingBox + """Detached obstacle world boxes; center and half-extents have shape (N, 3).""" - obstacle_max: torch.Tensor - """(B, 3) world-space max corner of the obstacle box.""" + tie_break_sign: float = 1.0 + """Directed sign used when coincident boxes have equally deep escape directions.""" -def compute_no_overlap_loss_aabb( +def compute_no_overlap_loss_obb( state: RelationSolverState, no_collision_strategy: NoCollisionLossStrategy, clearance_m: float, @@ -51,7 +46,7 @@ def compute_no_overlap_loss_aabb( skip_mesh_pairs: bool = False, debug: bool = False, ) -> tuple[torch.Tensor, int]: - """AABB collision loss summed over all directed pairs, returned per environment. + """OBB collision loss summed over all directed pairs, returned per environment. - Non-anchor vs fixed obstacle (anchor or collision object): gradient flows to the non-anchor only. - Non-anchor vs non-anchor: both objects accumulate gradient (two directed passes). @@ -83,23 +78,24 @@ def compute_no_overlap_loss_aabb( on_pairs.add((id(obj), id(rel.parent))) on_pairs.add((id(rel.parent), id(obj))) - extents: dict[PlaceableAsset | CollisionObject, tuple[torch.Tensor, torch.Tensor]] = {} + world_boxes: dict[PlaceableAsset | CollisionObject, OrientedBoundingBox] = {} for obj in non_anchor_objects: pos = state.get_position(obj) bbox = state.get_bbox(obj) - extents[obj] = (pos + bbox.min_point, pos + bbox.max_point) + world_boxes[obj] = bbox.translated(pos) for obstacle in fixed_obstacles: - obstacle_world_bbox = state.get_fixed_obstacle_world_bbox(obstacle) - extents[obstacle] = ( - obstacle_world_bbox.min_point.expand(batch_size, 3), - obstacle_world_bbox.max_point.expand(batch_size, 3), + bbox = state.get_fixed_obstacle_world_bbox(obstacle) + world_boxes[obstacle] = OrientedBoundingBox.from_tensors_unchecked( + bbox.center.expand(batch_size, 3), + bbox.half_extents.expand(batch_size, 3), + bbox.rotation_xyzw.expand(batch_size, 4), ) + detached_world_boxes = {obj: _detached(bbox) for obj, bbox in world_boxes.items()} pairs: list[NoOverlapPair] = [] pair_names: list[tuple[str, str]] = [] for child in non_anchor_objects: - child_min, child_max = extents[child] for obstacle in fixed_obstacles: if (id(child), id(obstacle)) in on_pairs: continue @@ -111,12 +107,10 @@ def compute_no_overlap_loss_aabb( ) ): continue - obstacle_min, obstacle_max = extents[obstacle] - pairs.append(NoOverlapPair(child_min, child_max, obstacle_min, obstacle_max)) + pairs.append(NoOverlapPair(world_boxes[child], detached_world_boxes[obstacle])) pair_names.append((child.name, obstacle.name)) for i, child in enumerate(non_anchor_objects): - child_min, child_max = extents[child] for j in range(i + 1, len(non_anchor_objects)): other = non_anchor_objects[j] if (id(child), id(other)) in on_pairs: @@ -129,24 +123,28 @@ def compute_no_overlap_loss_aabb( ) ): continue - other_min, other_max = extents[other] - pairs.append(NoOverlapPair(child_min, child_max, other_min.detach(), other_max.detach())) + pairs.append(NoOverlapPair(world_boxes[child], detached_world_boxes[other], tie_break_sign=1.0)) pair_names.append((child.name, other.name)) - pairs.append(NoOverlapPair(other_min, other_max, child_min.detach(), child_max.detach())) + pairs.append(NoOverlapPair(world_boxes[other], detached_world_boxes[child], tie_break_sign=-1.0)) pair_names.append((other.name, child.name)) num_pairs = len(pairs) if not pairs: return zero_loss, 0 - subject_min = torch.stack([p.subject_min for p in pairs], dim=0) - subject_max = torch.stack([p.subject_max for p in pairs], dim=0) - obstacle_min = torch.stack([p.obstacle_min for p in pairs], dim=0) - obstacle_max = torch.stack([p.obstacle_max for p in pairs], dim=0) - - pair_loss = no_collision_strategy.compute_loss_batched( - clearance_m, subject_min, subject_max, obstacle_min, obstacle_max - ) + subjects = _concatenate_boxes([pair.subject for pair in pairs]) + obstacles = _concatenate_boxes([pair.obstacle for pair in pairs]) + tie_break_signs = torch.tensor( + [pair.tie_break_sign for pair in pairs], + dtype=subjects.center.dtype, + device=device, + ).repeat_interleave(batch_size) + penetration = subjects.penetration( + obstacles, + clearance_m, + tie_break_sign=tie_break_signs, + ).reshape(num_pairs, batch_size) + pair_loss = no_collision_strategy.compute_loss_batched(penetration) if debug: for (subject_name, obstacle_name), loss in zip(pair_names, pair_loss): @@ -155,6 +153,24 @@ def compute_no_overlap_loss_aabb( return pair_loss.sum(dim=0), num_pairs +def _detached(bbox: OrientedBoundingBox) -> OrientedBoundingBox: + """Return a box detached from the optimization graph.""" + return OrientedBoundingBox.from_tensors_unchecked( + bbox.center.detach(), + bbox.half_extents.detach(), + bbox.rotation_xyzw.detach(), + ) + + +def _concatenate_boxes(boxes: list[OrientedBoundingBox]) -> OrientedBoundingBox: + """Concatenate pair-major box batches without revalidating their tensors.""" + return OrientedBoundingBox.from_tensors_unchecked( + torch.cat([bbox.center for bbox in boxes], dim=0), + torch.cat([bbox.half_extents for bbox in boxes], dim=0), + torch.cat([bbox.rotation_xyzw for bbox in boxes], dim=0), + ) + + def _fixed_pair_is_covered_by_mesh_collision( state: RelationSolverState, subject: PlaceableAsset, @@ -200,4 +216,4 @@ def _has_mesh_or_invariant_bbox( 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() + return state.get_base_bbox(obj).is_batch_invariant() diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py index 0035172d78..2ef9252514 100644 --- a/isaaclab_arena/relations/object_placer.py +++ b/isaaclab_arena/relations/object_placer.py @@ -5,6 +5,7 @@ from __future__ import annotations +import math import torch from dataclasses import dataclass, field from typing import TYPE_CHECKING @@ -23,10 +24,10 @@ get_anchor_objects, get_relation, ) -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose, PosePerEnv from isaaclab_arena.utils.random import get_random_rotation -from isaaclab_arena.utils.yaw import rotate_quat_by_yaw, wrap_angle_to_pi, yaw_from_quat_xyzw, yaw_toward_positions +from isaaclab_arena.utils.yaw import rotate_quat_by_yaw, yaw_toward_positions if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject @@ -47,8 +48,8 @@ class PlacementCandidate: validation_results: PlacementValidationResults """Per-check validation results for this candidate's layout.""" - orientations: dict[PlaceableAsset, float] = field(default_factory=dict) - """Placement-computed absolute world Z-yaws. Omitted objects retain their marker orientation.""" + rotations: dict[PlaceableAsset, tuple[float, float, float, float]] = field(default_factory=dict) + """Placement-computed xyzw rotations for movable objects.""" @property def is_valid(self) -> bool: @@ -128,8 +129,8 @@ def place( if self.params.apply_positions_to_objects: positions_per_env = [r.positions for r in results_per_env] - orientations_per_env = [r.orientations for r in results_per_env] - self._apply_poses(positions_per_env, anchor_objects_set, orientations_per_env) + rotations_per_env = [r.rotations for r in results_per_env] + self._apply_poses(positions_per_env, anchor_objects_set, rotations_per_env) return results_per_env @@ -198,10 +199,6 @@ def _prepare_placement( generator = torch.Generator() return set(anchor_objects), generator - # ------------------------------------------------------------------ - # Placement strategies - # ------------------------------------------------------------------ - def _place_ranked( self, objects: list[PlaceableAsset], @@ -223,41 +220,32 @@ def _place_ranked( assign_variants_for_envs(objects, num_envs, placement_seed=self.params.placement_seed) num_candidates = num_envs * candidates_per_env env_bboxes = build_per_env_bounding_boxes(objects, num_envs) - unrotated_candidate_bboxes = env_bboxes.get_bounding_boxes_for_solver_candidates(candidates_per_env) + candidate_bboxes = env_bboxes.get_bounding_boxes_for_solver_candidates(candidates_per_env) per_env_bboxes = env_bboxes.get_bounding_boxes_for_all_envs() initial_positions: list[dict[PlaceableAsset, tuple[float, float, float]]] = [] - orientations_per_candidate: list[dict[PlaceableAsset, float]] = [] + rotations_per_candidate: list[dict[PlaceableAsset, tuple[float, float, float, float]]] = [] for candidate_idx in range(num_candidates): cur_env = candidate_idx // candidates_per_env if generator is not None: assert self.params.placement_seed is not None generator.manual_seed(self.params.placement_seed + candidate_idx) + rotations = self._generate_initial_rotations(objects, anchor_objects_set, generator) + rotations_per_candidate.append(rotations) initial_positions.append( - self._generate_initial_positions(objects, anchor_objects_set, per_env_bboxes[cur_env], generator) - ) - orientations_per_candidate.append( - self._generate_initial_orientations(objects, anchor_objects_set, generator) + self._generate_initial_positions( + objects, anchor_objects_set, per_env_bboxes[cur_env], generator, rotations=rotations + ) ) - # Bake each candidate's yaw into a conservative enclosing bbox for overlap checks. - candidate_bboxes = self._rotate_candidate_bboxes( - objects, unrotated_candidate_bboxes, orientations_per_candidate - ) - all_positions = self._solver.solve( objects, initial_positions, env_bboxes=candidate_bboxes, - env_bboxes_include_yaw=any(orientations for orientations in orientations_per_candidate), - orientations=orientations_per_candidate, + rotations=rotations_per_candidate, collision_objects=collision_objects, ) - self._apply_face_to_orientations(all_positions, orientations_per_candidate) - # FaceTo yaw is only known after solving, so rebuild from unrotated boxes before validation. - candidate_bboxes = self._rotate_candidate_bboxes( - objects, unrotated_candidate_bboxes, orientations_per_candidate - ) + self._apply_face_to_rotations(all_positions, rotations_per_candidate) assert self._solver.last_loss_per_env is not None all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist() bboxes_per_candidate = [ @@ -265,19 +253,19 @@ def _place_ranked( for candidate_idx in range(num_candidates) ] all_validations = self._validate_candidates( - all_positions, orientations_per_candidate, bboxes_per_candidate, collision_objects + all_positions, rotations_per_candidate, bboxes_per_candidate, collision_objects ) - candidates: list[PlacementCandidate] = [] - for candidate_idx in range(num_candidates): - candidates.append( - PlacementCandidate( - all_losses[candidate_idx], - all_positions[candidate_idx], - all_validations[candidate_idx], - orientations_per_candidate[candidate_idx], - ) + candidates = [ + PlacementCandidate(loss, position, validation, rotation) + for loss, position, validation, rotation in zip( + all_losses, + all_positions, + all_validations, + rotations_per_candidate, + strict=True, ) + ] ranked_candidate_slices = self._rank_candidates(candidates, num_envs, candidates_per_env) ranked_results = [ @@ -287,7 +275,7 @@ def _place_ranked( positions=candidate.positions, final_loss=candidate.loss, attempts=attempts_per_result, - orientations=candidate.orientations, + rotations=candidate.rotations, ) for candidate in candidate_slice ] @@ -334,8 +322,9 @@ def _generate_initial_positions( self, objects: list[PlaceableAsset], anchor_objects: set[PlaceableAsset], - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox], generator: torch.Generator | None = None, + rotations: dict[PlaceableAsset, tuple[float, float, float, float]] | None = None, ) -> dict[PlaceableAsset, tuple[float, float, float]]: """Generate initial positions for all objects. @@ -344,7 +333,7 @@ def _generate_initial_positions( anchor's center; the solver handles their placement from there. Args: - env_bboxes: Per-object bboxes for the current env, each with shape (1, 3). + env_bboxes: Per-object bounding boxes for the current environment, each with N=1. generator: Optional RNG generator for reproducible sampling. When None, uses PyTorch's global RNG. @@ -355,6 +344,14 @@ def _generate_initial_positions( anchor_bbox = self._get_world_bbox_for_init(first_anchor, env_bboxes) cx, cy, cz = float(anchor_bbox.center[0, 0]), float(anchor_bbox.center[0, 1]), float(anchor_bbox.center[0, 2]) + candidate_bboxes = { + obj: ( + bbox.rotated_by_quat(rotations[obj]) + if rotations is not None and obj in rotations and obj not in anchor_objects + else bbox + ) + for obj, bbox in env_bboxes.items() + } positions: dict[PlaceableAsset, tuple[float, float, float]] = {} for obj in objects: @@ -367,7 +364,7 @@ def _generate_initial_positions( positions[obj] = initial_pose.position_xyz elif any(isinstance(r, On) for r in obj.get_relations()): positions[obj] = self._compute_on_guided_position( - obj, anchor_objects, anchor_bbox, env_bboxes, generator + obj, anchor_objects, anchor_bbox, candidate_bboxes, generator ) else: positions[obj] = (cx, cy, cz) @@ -376,55 +373,44 @@ def _generate_initial_positions( @staticmethod def _get_world_bbox_for_init( obj: PlaceableAsset, - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], - ) -> AxisAlignedBoundingBox: + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox], + ) -> OrientedBoundingBox: initial_pose = obj.get_initial_pose() assert isinstance( initial_pose, Pose ), f"Object '{obj.name}' must have a fixed Pose to use its env bbox, got {type(initial_pose).__name__}." - return env_bboxes[obj].translated(initial_pose.position_xyz) + return env_bboxes[obj].transformed(initial_pose.position_xyz, initial_pose.rotation_xyzw) - def _generate_initial_orientations( + def _generate_initial_rotations( self, objects: list[PlaceableAsset], anchor_objects: set[PlaceableAsset], generator: torch.Generator | None = None, - ) -> dict[PlaceableAsset, float]: - """Sample absolute world Z-yaws for non-anchor objects without FaceTo. - - Marker yaw is included; random_yaw_init adds a sampled delta. Roll/pitch marker objects are - omitted so their requested rotation is applied verbatim; their footprint is enclosed by - _rotate_candidate_bboxes so overlap validation stays sound. - """ - orientations: dict[PlaceableAsset, float] = {} + ) -> dict[PlaceableAsset, tuple[float, float, float, float]]: + """Return fixed candidate xyzw rotations for non-FaceTo movable objects.""" + rotations: dict[PlaceableAsset, tuple[float, float, float, float]] = {} for obj in objects: marker = get_relation(obj, RotateAroundSolution) - has_roll_pitch = marker is not None and (marker.roll_rad != 0.0 or marker.pitch_rad != 0.0) - marker_yaw = marker.yaw_rad if marker is not None else 0.0 if obj in anchor_objects: - assert marker is None or (marker_yaw == 0.0 and not has_roll_pitch), ( + assert marker is None, ( f"Anchor '{obj.name}' has a RotateAroundSolution. " "Anchors are not repositioned by the placer, so any marker rotation must " "already be baked into the anchor's initial_pose before calling place()." ) - elif get_relation(obj, FaceTo) is None and not has_roll_pitch: + elif get_relation(obj, FaceTo) is None: + base = marker.get_rotation_xyzw() if marker is not None else (0.0, 0.0, 0.0, 1.0) sampled_yaw = get_random_rotation(generator) if self.params.random_yaw_init else 0.0 - total_yaw = wrap_angle_to_pi(sampled_yaw + marker_yaw) - if total_yaw != 0.0: - orientations[obj] = total_yaw - return orientations + rotations[obj] = rotate_quat_by_yaw(base, sampled_yaw) + return rotations @staticmethod - def _apply_face_to_orientations( + def _apply_face_to_rotations( positions_per_candidate: list[dict[PlaceableAsset, tuple[float, float, float]]], - orientations_per_candidate: list[dict[PlaceableAsset, float]], + rotations_per_candidate: list[dict[PlaceableAsset, tuple[float, float, float, float]]], ) -> None: - """Write defined FaceTo yaws into each candidate's orientation dictionary in place. - - Undefined directions leave the subject absent from the dictionary. - """ + """Store final FaceTo rotations for validation and pose application.""" assert positions_per_candidate, "positions_per_candidate must not be empty" - assert len(positions_per_candidate) == len(orientations_per_candidate) + assert len(positions_per_candidate) == len(rotations_per_candidate) objects = positions_per_candidate[0] for obj in objects: relation = get_relation(obj, FaceTo) @@ -435,57 +421,29 @@ def _apply_face_to_orientations( yaws, is_defined = yaw_toward_positions(subject_positions, target_positions) for candidate_idx, (yaw, direction_is_defined) in enumerate(zip(yaws, is_defined, strict=True)): if direction_is_defined: - orientations_per_candidate[candidate_idx][obj] = yaw.item() - - @staticmethod - def _rotate_candidate_bboxes( - objects: list[PlaceableAsset], - candidate_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], - orientations_per_candidate: list[dict[PlaceableAsset, float]], - ) -> dict[PlaceableAsset, AxisAlignedBoundingBox]: - """Replace each candidate's bbox with the AABB enclosing its fully-oriented object. - - Composes each object's static RotateAroundSolution marker rotation (roll/pitch/yaw) with the - per-candidate yaw (sampled + FaceTo, carried as absolute world yaw in orientations_per_candidate) - and refits the box to that combined quaternion -- the same composition _apply_poses uses for the - final pose, so overlap boxes match the placed object regardless of rotation axis. Objects with - no rotation are returned unchanged, keeping the no-rotation path exact. - """ - num_candidates = len(orientations_per_candidate) - rotated: dict[PlaceableAsset, AxisAlignedBoundingBox] = {} - for obj in objects: - bbox = candidate_bboxes[obj] - marker = get_relation(obj, RotateAroundSolution) - marker_rotation = marker.get_rotation_xyzw() if marker is not None else (0.0, 0.0, 0.0, 1.0) - has_roll_pitch = marker is not None and (marker.roll_rad != 0.0 or marker.pitch_rad != 0.0) - # orientations carries absolute world yaw; subtract the marker's own yaw to get the delta to compose. - marker_yaw = yaw_from_quat_xyzw(marker_rotation) - extra_yaws = [ - orientations_per_candidate[c].get(obj, marker_yaw) - marker_yaw for c in range(num_candidates) - ] - if not has_roll_pitch and all(yaw == 0.0 for yaw in extra_yaws): - rotated[obj] = bbox - else: - quats = [rotate_quat_by_yaw(marker_rotation, yaw) for yaw in extra_yaws] - quat_tensor = torch.tensor(quats, dtype=torch.float32, device=bbox.min_point.device) - rotated[obj] = bbox.rotated_by_quat(quat_tensor) - return rotated + half_yaw = yaw.item() * 0.5 + rotations_per_candidate[candidate_idx][obj] = ( + 0.0, + 0.0, + math.sin(half_yaw), + math.cos(half_yaw), + ) @staticmethod def _get_bounding_boxes_for_candidate_index( - bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], + bboxes: dict[PlaceableAsset, OrientedBoundingBox], candidate_idx: int, - ) -> dict[PlaceableAsset, AxisAlignedBoundingBox]: - """Slice one candidate's bboxes (each (1, 3)) out of the stacked (num_candidates, 3) boxes.""" + ) -> dict[PlaceableAsset, OrientedBoundingBox]: + """Return one candidate's bounding boxes, each with N=1.""" return {obj: bbox[candidate_idx] for obj, bbox in bboxes.items()} def _get_on_parent_world_bbox( self, parent: PlaceableAsset, anchor_objects: set[PlaceableAsset], - anchor_bbox: AxisAlignedBoundingBox, - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], - ) -> AxisAlignedBoundingBox: + anchor_bbox: OrientedBoundingBox, + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox], + ) -> OrientedBoundingBox: """Resolve the world bbox of an On relation's parent for initialization purposes. If the parent is an anchor, return its world bbox directly. @@ -506,8 +464,8 @@ def _compute_on_guided_position( self, obj: PlaceableAsset, anchor_objects: set[PlaceableAsset], - anchor_bbox: AxisAlignedBoundingBox, - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], + anchor_bbox: OrientedBoundingBox, + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox], generator: torch.Generator | None = None, ) -> tuple[float, float, float]: """Compute an initial position for an object with an On relation. @@ -516,31 +474,38 @@ def _compute_on_guided_position( so the solver starts from a valid region. Args: - env_bboxes: Per-object bboxes for the current env, each with shape (1, 3). + env_bboxes: Per-object bounding boxes for the current environment, each with N=1. generator: Optional RNG generator for reproducible sampling. When None, uses PyTorch's global RNG. """ on_relation = next(r for r in obj.get_relations() if isinstance(r, On)) parent_bbox = self._get_on_parent_world_bbox(on_relation.parent, anchor_objects, anchor_bbox, env_bboxes) child_bbox = env_bboxes[obj] + axes = torch.eye(3, dtype=child_bbox.center.dtype, device=child_bbox.center.device) + parent_x_min, parent_x_max = parent_bbox.get_bounds_along_axis(axes[0]) + parent_y_min, parent_y_max = parent_bbox.get_bounds_along_axis(axes[1]) + _, parent_z_max = parent_bbox.get_bounds_along_axis(axes[2]) + child_x_min, child_x_max = child_bbox.get_bounds_along_axis(axes[0]) + child_y_min, child_y_max = child_bbox.get_bounds_along_axis(axes[1]) + child_z_min, _ = child_bbox.get_bounds_along_axis(axes[2]) x = self._sample_axis_position( - parent_bbox.min_point[0, 0], - parent_bbox.max_point[0, 0], - child_bbox.min_point[0, 0], - child_bbox.max_point[0, 0], + parent_x_min[0], + parent_x_max[0], + child_x_min[0], + child_x_max[0], generator, ) y = self._sample_axis_position( - parent_bbox.min_point[0, 1], - parent_bbox.max_point[0, 1], - child_bbox.min_point[0, 1], - child_bbox.max_point[0, 1], + parent_y_min[0], + parent_y_max[0], + child_y_min[0], + child_y_max[0], generator, ) # Convert from child-origin Z to child-bottom Z so the bottom face lands on the parent top. - z = float(parent_bbox.max_point[0, 2] + on_relation.clearance_m - child_bbox.min_point[0, 2]) + z = float(parent_z_max[0] + on_relation.clearance_m - child_z_min[0]) return (x, y, z) @@ -577,8 +542,8 @@ def _sample_axis_position( def _validate_candidates( self, positions: list[dict[PlaceableAsset, tuple[float, float, float]]], - orientations: list[dict[PlaceableAsset, float]], - bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]], + bboxes: list[dict[PlaceableAsset, OrientedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[PlacementValidationResults]: """Run every enabled validator over all candidates and collect per-candidate results. @@ -588,8 +553,8 @@ def _validate_candidates( Args: positions: Solved (x, y, z) per object, one dict per candidate. - orientations: Absolute world Z-yaw per object, one dict per candidate (may be empty). - bboxes: Per-object bboxes for each candidate's env, each (1, 3). + rotations: Candidate xyzw rotation per movable object. + bboxes: Per-object bounding boxes for each candidate's environment, each with N=1. collision_objects: Fixed background obstacles shared across candidates. """ # required_checks=None means "every enabled check is required"; an empty set means no checks. @@ -601,7 +566,7 @@ def _validate_candidates( self._run_inexpensive_checks( positions, - orientations, + rotations, bboxes, collision_objects, layout_pass_verdicts_by_check, @@ -609,7 +574,7 @@ def _validate_candidates( ) self._run_expensive_checks( positions, - orientations, + rotations, bboxes, collision_objects, required, @@ -635,8 +600,8 @@ def _validate_candidates( def _run_inexpensive_checks( self, positions: list[dict[PlaceableAsset, tuple[float, float, float]]], - orientations: list[dict[PlaceableAsset, float]], - bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]], + bboxes: list[dict[PlaceableAsset, OrientedBoundingBox]], collision_objects: list[CollisionObject], layout_pass_verdicts_by_check: dict[str, list[bool]], num_layouts_evaluated_by_check: dict[str, int], @@ -646,15 +611,15 @@ def _run_inexpensive_checks( for validator in self._validators: if not validator.run_after_inexpensive_checks: layout_pass_verdicts_by_check[validator.check] = validator.validate_batch( - positions, orientations, bboxes, collision_objects + positions, rotations, bboxes, collision_objects ) num_layouts_evaluated_by_check[validator.check] = num_candidates def _run_expensive_checks( self, positions: list[dict[PlaceableAsset, tuple[float, float, float]]], - orientations: list[dict[PlaceableAsset, float]], - bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]], + bboxes: list[dict[PlaceableAsset, OrientedBoundingBox]], collision_objects: list[CollisionObject], required: set[str] | None, layout_pass_verdicts_by_check: dict[str, list[bool]], @@ -672,7 +637,7 @@ def _run_expensive_checks( # only passed layouts are validated verdicts_over_passed_layout = validator.validate_batch( [positions[i] for i in passed_layout_indices], - [orientations[i] for i in passed_layout_indices], + [rotations[i] for i in passed_layout_indices], [bboxes[i] for i in passed_layout_indices], collision_objects, ) @@ -702,29 +667,18 @@ def _apply_poses( self, positions_per_env: list[dict[PlaceableAsset, tuple[float, float, float]]], anchor_objects: set[PlaceableAsset], - orientations_per_env: list[dict[PlaceableAsset, float]], + rotations_per_env: list[dict[PlaceableAsset, tuple[float, float, float, float]]], ) -> None: - """Apply solved positions and orientations to non-anchor objects. - - orientations_per_env carries absolute world yaw; marker yaw is subtracted before composition. - """ + """Apply solved positions and xyzw rotations to non-anchor objects.""" num_envs = len(positions_per_env) objects = list(positions_per_env[0]) for obj in objects: if obj in anchor_objects: continue - rotate_marker = get_relation(obj, RotateAroundSolution) - marker_rotation = rotate_marker.get_rotation_xyzw() if rotate_marker else (0.0, 0.0, 0.0, 1.0) - marker_yaw = yaw_from_quat_xyzw(marker_rotation) - - def _yaw_delta(env_idx: int) -> float: - """Return the yaw to compose with the RotateAroundSolution marker rotation.""" - return orientations_per_env[env_idx].get(obj, marker_yaw) - marker_yaw - if num_envs == 1: pos = positions_per_env[0][obj] - rotation_xyzw = rotate_quat_by_yaw(marker_rotation, _yaw_delta(0)) + rotation_xyzw = rotations_per_env[0].get(obj, (0.0, 0.0, 0.0, 1.0)) random_marker = get_relation(obj, RandomAroundSolution) if random_marker is not None: obj.set_initial_pose(random_marker.to_pose_range_centered_at(pos, rotation_xyzw=rotation_xyzw)) @@ -734,7 +688,7 @@ def _yaw_delta(env_idx: int) -> float: poses = [ Pose( position_xyz=positions_per_env[env_idx][obj], - rotation_xyzw=rotate_quat_by_yaw(marker_rotation, _yaw_delta(env_idx)), + rotation_xyzw=rotations_per_env[env_idx].get(obj, (0.0, 0.0, 0.0, 1.0)), ) for env_idx in range(num_envs) ] diff --git a/isaaclab_arena/relations/placement_asset.py b/isaaclab_arena/relations/placement_asset.py index 2478a4ba7a..4a6a216b5d 100644 --- a/isaaclab_arena/relations/placement_asset.py +++ b/isaaclab_arena/relations/placement_asset.py @@ -13,7 +13,7 @@ from isaaclab_arena.assets.asset import Asset from isaaclab_arena.relations.collision_mode import CollisionMode from isaaclab_arena.relations.relations import IsAnchor, Relation, RelationBase, RequiresReachability, UnaryRelation -from isaaclab_arena.utils.bounding_box import quaternion_to_90_deg_z_quarters +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange if TYPE_CHECKING: @@ -21,8 +21,6 @@ from isaaclab.managers import EventTermCfg - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox - class PlaceableAsset(Asset, ABC): """Asset whose root pose can be constrained by spatial relations.""" @@ -135,11 +133,11 @@ def has_unplaced_auxiliary_prims(self) -> bool: return False @abstractmethod - def get_bounding_box(self) -> AxisAlignedBoundingBox: - """Return root-relative axis-aligned bounds.""" + def get_bounding_box(self) -> OrientedBoundingBox: + """Return root-relative oriented bounds.""" - def get_world_bounding_box(self) -> AxisAlignedBoundingBox: - """Return bounds transformed by a fixed root pose with a quarter-turn Z rotation. + def get_world_bounding_box(self) -> OrientedBoundingBox: + """Return bounds transformed by a fixed root pose. Unset, ranged, and per-environment poses leave the root-relative bounds unchanged. """ @@ -147,11 +145,10 @@ def get_world_bounding_box(self) -> AxisAlignedBoundingBox: initial_pose = self.get_initial_pose() if not isinstance(initial_pose, Pose): return bounding_box - quarters = quaternion_to_90_deg_z_quarters(initial_pose.rotation_xyzw) - return bounding_box.rotated_90_around_z(quarters).translated(initial_pose.position_xyz) + return bounding_box.transformed(initial_pose.position_xyz, initial_pose.rotation_xyzw) def get_collision_mesh(self) -> trimesh.Trimesh | None: - """Return this asset's collision mesh, or ``None`` to fall back to the axis-aligned bounds. + """Return this asset's collision mesh, or ``None`` to fall back to its bounds. Concrete (not abstract) so assets without a mesh simply keep the ``None`` default. """ diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py index bed1ca3d7e..491b92b278 100644 --- a/isaaclab_arena/relations/placement_events.py +++ b/isaaclab_arena/relations/placement_events.py @@ -11,7 +11,6 @@ from isaaclab_arena.relations.relations import RotateAroundSolution, get_anchor_objects from isaaclab_arena.utils.pose import Pose from isaaclab_arena.utils.velocity import Velocity -from isaaclab_arena.utils.yaw import rotate_quat_by_yaw, yaw_from_quat_xyzw if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -58,10 +57,7 @@ def get_base_rotation_per_asset( def get_pose_from_layout(asset: PlaceableAsset, layout: PlacementResult) -> Pose: """Return an asset pose from a solved layout.""" assert asset in layout.positions, f"Placement layout is missing non-anchor asset '{asset.name}'" - base_rotation = get_rotation_xyzw(asset) - marker_yaw = yaw_from_quat_xyzw(base_rotation) - total_yaw = layout.orientations.get(asset, marker_yaw) - rotation = rotate_quat_by_yaw(base_rotation, total_yaw - marker_yaw) + rotation = layout.rotations.get(asset, IDENTITY_ROTATION_XYZW) return Pose(position_xyz=layout.positions[asset], rotation_xyzw=rotation) diff --git a/isaaclab_arena/relations/placement_result.py b/isaaclab_arena/relations/placement_result.py index 84bdc2da03..1bace4c46a 100644 --- a/isaaclab_arena/relations/placement_result.py +++ b/isaaclab_arena/relations/placement_result.py @@ -29,8 +29,8 @@ class PlacementResult: attempts: int """Number of attempts made.""" - orientations: dict[PlaceableAsset, float] = field(default_factory=dict) - """Sparse map of world yaw angles ``theta_z`` in radians; omitted assets retain marker orientation.""" + rotations: dict[PlaceableAsset, tuple[float, float, float, float]] = field(default_factory=dict) + """Sparse movable-asset world rotations in xyzw order; omitted assets use identity.""" @property def success(self) -> bool: diff --git a/isaaclab_arena/relations/placement_validators.py b/isaaclab_arena/relations/placement_validators.py index 65ee5612eb..7f65810fa8 100644 --- a/isaaclab_arena/relations/placement_validators.py +++ b/isaaclab_arena/relations/placement_validators.py @@ -5,6 +5,7 @@ from __future__ import annotations +import numpy as np import torch import trimesh from abc import ABC, abstractmethod @@ -12,6 +13,7 @@ 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.no_overlap_mesh import transform_points_between_frames 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,22 +24,19 @@ ) 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 OrientedBoundingBox from isaaclab_arena.utils.pose import Pose -from isaaclab_arena.utils.yaw import centers_in_target_frame, yaw_from_quat_xyzw, yaw_toward_positions +from isaaclab_arena.utils.yaw import yaw_toward_positions if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject 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): - """A single build-time placement check evaluated over a batch of candidate layouts. - - Register a concrete validator with @register_validator so build_validators() can discover it. - """ + """A build-time placement check over a batch of candidate layouts.""" check: ClassVar[str] """The check name this validator reports; its registry key and result key. Built-ins use a @@ -63,16 +62,16 @@ def is_available(cls, params: ObjectPlacerParams) -> bool: def validate_batch( self, positions: list[dict[PlaceableAsset, tuple[float, float, float]]], - orientations: list[dict[PlaceableAsset, float]], - bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]], + bboxes: list[dict[PlaceableAsset, OrientedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: """Return one pass/fail verdict per candidate layout. Args: positions: Solved (x, y, z) per object, one dict per candidate. - orientations: Absolute world Z-yaw per object, one dict per candidate (may be empty). - bboxes: Per-object bboxes for the candidate's env, one dict per candidate, each (1, 3). + rotations: World xyzw rotations per movable object. + bboxes: Per-object boxes for each candidate's environment, each with N=1. collision_objects: Fixed background obstacles shared across candidates. """ pass @@ -108,6 +107,31 @@ def build_validators(params: ObjectPlacerParams) -> list[PlacementValidator]: return validators +def _candidate_local_bbox( + obj: PlaceableAsset, + bbox: OrientedBoundingBox, + rotations: dict[PlaceableAsset, tuple[float, float, float, float]], +) -> OrientedBoundingBox: + """Return candidate-oriented local geometry for a movable object.""" + if obj.is_anchor: + return bbox + return bbox.rotated_by_quat(rotations.get(obj, (0.0, 0.0, 0.0, 1.0))) + + +def _candidate_world_bbox( + obj: PlaceableAsset, + bbox: OrientedBoundingBox, + position: tuple[float, float, float], + rotations: dict[PlaceableAsset, tuple[float, float, float, float]], +) -> OrientedBoundingBox: + """Transform candidate or fixed-anchor geometry into world coordinates.""" + if obj.is_anchor: + pose = obj.get_initial_pose() + assert isinstance(pose, Pose), f"Anchor '{obj.name}' must have a fixed Pose." + return bbox.transformed(pose.position_xyz, pose.rotation_xyzw) + return bbox.transformed(position, rotations.get(obj, (0.0, 0.0, 0.0, 1.0))) + + @register_validator class OnRelationValidator(PlacementValidator): """Validate every On relation: child rests on its parent within X/Y footprint and Z band.""" @@ -117,16 +141,17 @@ class OnRelationValidator(PlacementValidator): def validate_batch( self, positions: list[dict[PlaceableAsset, tuple[float, float, float]]], - orientations: list[dict[PlaceableAsset, float]], - bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]], + bboxes: list[dict[PlaceableAsset, OrientedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: - return [self._validate(positions[i], bboxes[i]) for i in range(len(positions))] + return [self._validate(positions[i], bboxes[i], rotations[i]) for i in range(len(positions))] def _validate( self, positions: dict[PlaceableAsset, tuple[float, float, float]], - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox], + rotations: dict[PlaceableAsset, tuple[float, float, float, float]] | None = None, ) -> bool: """Validate each On relation; keep in sync with OnLossStrategy in relation_loss_strategies.py. @@ -136,8 +161,9 @@ def _validate( Args: positions: Solved positions for each object. - env_bboxes: Per-object bboxes for the current env, each with shape (1, 3). + env_bboxes: Per-object boxes for the current environment, each with N=1. """ + rotations = rotations or {} for obj in positions: for rel in obj.get_relations(): if not isinstance(rel, On): @@ -145,21 +171,27 @@ def _validate( parent = rel.parent if parent not in positions: continue - child_bbox = env_bboxes[obj] - parent_bbox = env_bboxes[parent] - child_world = child_bbox.translated(positions[obj]) - parent_world = parent_bbox.translated(positions[parent]) - parent_size = parent_world.max_point - parent_world.min_point - child_size = child_world.max_point - child_world.min_point + child_bbox = _candidate_local_bbox(obj, env_bboxes[obj], rotations) + child_world = _candidate_world_bbox(obj, env_bboxes[obj], positions[obj], rotations) + parent_world = _candidate_world_bbox(parent, env_bboxes[parent], positions[parent], rotations) + axes = torch.eye(3, dtype=child_bbox.center.dtype, device=child_bbox.center.device) + child_x_min, child_x_max = child_world.get_bounds_along_axis(axes[0]) + child_y_min, child_y_max = child_world.get_bounds_along_axis(axes[1]) + child_z_min, _ = child_world.get_bounds_along_axis(axes[2]) + parent_x_min, parent_x_max = parent_world.get_bounds_along_axis(axes[0]) + parent_y_min, parent_y_max = parent_world.get_bounds_along_axis(axes[1]) + _, parent_z_max = parent_world.get_bounds_along_axis(axes[2]) + parent_size = torch.stack([parent_x_max - parent_x_min, parent_y_max - parent_y_min], dim=1) + child_size = torch.stack([child_x_max - child_x_min, child_y_max - child_y_min], dim=1) m = rel.edge_margin_m # 1) Checking that with the specified margin, the parent is wide enough to place the child on top if m > 0.0: freespace = parent_size - child_size # A margin too large for the surface inverts the inset band so containment can never pass. - if torch.any(freespace[0, :2] < 2 * m): + if torch.any(freespace[0] < 2 * m): # The maximum feasible margin is the minimum of the freespace on the xy axes. - max_feasible_margin = max(0.0, min(freespace[0, :2]) / 2.0) + max_feasible_margin = max(0.0, min(freespace[0]) / 2.0) # When parent < child, freespace[0, :2] is negative and max_feasible_margin is 0.0. if max_feasible_margin > 0.0: if self._params.verbose: @@ -171,22 +203,22 @@ def _validate( return False # 2) Checking that the child lies within the parent's xy if ( - child_world.min_point[0, 0] < parent_world.min_point[0, 0] + m - or child_world.max_point[0, 0] > parent_world.max_point[0, 0] - m - or child_world.min_point[0, 1] < parent_world.min_point[0, 1] + m - or child_world.max_point[0, 1] > parent_world.max_point[0, 1] - m + child_x_min[0] < parent_x_min[0] + m + or child_x_max[0] > parent_x_max[0] - m + or child_y_min[0] < parent_y_min[0] + m + or child_y_max[0] > parent_y_max[0] - m ): if self._params.verbose: print(f"On relation: '{obj.name}' XY outside parent (retrying)") return False # 3) Checking that the child lies within an acceptable z-range. - parent_local_top_z: float = parent_bbox.max_point[0, 2].item() - child_local_bottom_z: float = child_bbox.min_point[0, 2].item() - parent_top_z = parent_local_top_z + positions[parent][2] clearance_m = rel.clearance_m - child_bottom_z = child_local_bottom_z + positions[obj][2] eps_z = self._params.on_relation_z_tolerance_m - if child_bottom_z <= parent_top_z - eps_z or child_bottom_z > parent_top_z + clearance_m + eps_z: + numerical_eps = 1e-6 + if ( + child_z_min[0] <= parent_z_max[0] - eps_z + numerical_eps + or child_z_min[0] > parent_z_max[0] + clearance_m + eps_z + ): if self._params.verbose: print(f" On relation: '{obj.name}' Z outside band (retrying)") return False @@ -202,16 +234,17 @@ class NextToValidator(PlacementValidator): def validate_batch( self, positions: list[dict[PlaceableAsset, tuple[float, float, float]]], - orientations: list[dict[PlaceableAsset, float]], - bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]], + bboxes: list[dict[PlaceableAsset, OrientedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: - return [self._validate(positions[i], bboxes[i]) for i in range(len(positions))] + return [self._validate(positions[i], bboxes[i], rotations[i]) for i in range(len(positions))] def _validate( self, positions: dict[PlaceableAsset, tuple[float, float, float]], - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox], + rotations: dict[PlaceableAsset, tuple[float, float, float, float]] | None = None, ) -> bool: """Validate each NextTo relation: child on the requested side, facing edge within the relation's tolerance_m of distance_m from the parent edge. Shares next_to_violations with @@ -219,8 +252,9 @@ def _validate( Args: positions: Solved positions for each object. - env_bboxes: Per-object bboxes for the current env, each with shape (1, 3). + env_bboxes: Per-object boxes for the current environment, each with N=1. """ + rotations = rotations or {} for obj in positions: for rel in obj.get_relations(): if not isinstance(rel, NextTo): @@ -229,9 +263,9 @@ def _validate( if parent not in positions: continue cfg = SIDE_CONFIGS[rel.side] - child_bbox = env_bboxes[obj] - child_pos = child_bbox.min_point.new_tensor([positions[obj]]) - parent_world = env_bboxes[parent].translated(positions[parent]) + child_bbox = _candidate_local_bbox(obj, env_bboxes[obj], rotations) + child_pos = child_bbox.center.new_tensor([positions[obj]]) + parent_world = _candidate_world_bbox(parent, env_bboxes[parent], positions[parent], rotations) half_plane, distance = next_to_violations(cfg, child_pos, child_bbox, parent_world, rel.distance_m) if half_plane.item() > rel.tolerance_m or distance.item() > rel.tolerance_m: @@ -254,16 +288,17 @@ class NotNextToValidator(PlacementValidator): def validate_batch( self, positions: list[dict[PlaceableAsset, tuple[float, float, float]]], - orientations: list[dict[PlaceableAsset, float]], - bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]], + bboxes: list[dict[PlaceableAsset, OrientedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: - return [self._validate(positions[i], bboxes[i]) for i in range(len(positions))] + return [self._validate(positions[i], bboxes[i], rotations[i]) for i in range(len(positions))] def _validate( self, positions: dict[PlaceableAsset, tuple[float, float, float]], - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox], + rotations: dict[PlaceableAsset, tuple[float, float, float, float]] | None = None, ) -> bool: """Validate each NotNextTo relation: child has cleared the keep-out zone beside the parent (within the relation's tolerance_m) via either route — back over the edge or past the @@ -271,8 +306,9 @@ def _validate( Args: positions: Solved positions for each object. - env_bboxes: Per-object bboxes for the current env, each with shape (1, 3). + env_bboxes: Per-object boxes for the current environment, each with N=1. """ + rotations = rotations or {} for obj in positions: for rel in obj.get_relations(): if not isinstance(rel, NotNextTo): @@ -282,9 +318,9 @@ def _validate( continue cfg = SIDE_CONFIGS[rel.side] margin_m = self._not_next_to_margin(rel) - child_bbox = env_bboxes[obj] - child_pos = child_bbox.min_point.new_tensor([positions[obj]]) - parent_world = env_bboxes[parent].translated(positions[parent]) + child_bbox = _candidate_local_bbox(obj, env_bboxes[obj], rotations) + child_pos = child_bbox.center.new_tensor([positions[obj]]) + parent_world = _candidate_world_bbox(parent, env_bboxes[parent], positions[parent], rotations) remaining_side, remaining_cross = not_next_to_violations( cfg, child_pos, child_bbox, parent_world, margin_m ) @@ -315,16 +351,16 @@ class FaceToValidator(PlacementValidator): def validate_batch( self, positions: list[dict[PlaceableAsset, tuple[float, float, float]]], - orientations: list[dict[PlaceableAsset, float]], - bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]], + bboxes: list[dict[PlaceableAsset, OrientedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: - return [self._validate(positions[i], orientations[i]) for i in range(len(positions))] + return [self._validate(positions[i], rotations[i]) for i in range(len(positions))] def _validate( self, positions: dict[PlaceableAsset, tuple[float, float, float]], - orientations: dict[PlaceableAsset, float] | None, + rotations: dict[PlaceableAsset, tuple[float, float, float, float]] | None, ) -> bool: """Validate that every FaceTo subject has a defined direction and computed yaw.""" for obj in positions: @@ -338,7 +374,7 @@ def _validate( if self._params.verbose: print(f" FaceTo: '{obj.name}' is too close to its target in XY") return False - if orientations is None or obj not in orientations: + if rotations is None or obj not in rotations: if self._params.verbose: print(f" FaceTo: '{obj.name}' has no computed facing yaw") return False @@ -362,31 +398,30 @@ def __init__(self, params: ObjectPlacerParams) -> None: def validate_batch( self, positions: list[dict[PlaceableAsset, tuple[float, float, float]]], - orientations: list[dict[PlaceableAsset, float]], - bboxes: list[dict[PlaceableAsset, AxisAlignedBoundingBox]], + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]], + bboxes: list[dict[PlaceableAsset, OrientedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: - return [ - self._validate(positions[i], bboxes[i], orientations[i], collision_objects) for i in range(len(positions)) - ] + return [self._validate(positions[i], bboxes[i], rotations[i], collision_objects) for i in range(len(positions))] def _validate( self, positions: dict[PlaceableAsset, tuple[float, float, float]], - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], - orientations: dict[PlaceableAsset, float] | None, + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox], + rotations: dict[PlaceableAsset, tuple[float, float, float, float]], collision_objects: list[CollisionObject] | None, ) -> bool: - """AABB overlap check, falling through to mesh penetration for mesh-collision objects.""" + """OBB overlap check, falling through to mesh penetration when requested.""" use_mesh = self._should_validate_mesh(positions, collision_objects) no_overlap = self._validate_no_overlap( positions, env_bboxes, + rotations, collision_objects=collision_objects, skip_mesh_pairs=use_mesh, ) 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, rotations, collision_objects) return no_overlap def _should_validate_mesh( @@ -404,23 +439,18 @@ def _should_validate_mesh( @staticmethod def _collect_skip_pairs( positions: dict[PlaceableAsset, tuple[float, float, float]], - ) -> tuple[set[tuple], set[int]]: - """Build On-pair skip set and anchor ID set from positioned objects. - - Returns: - Tuple of (on_pairs, anchor_ids) where on_pairs contains (id(a), id(b)) - tuples for On-linked objects, and anchor_ids contains id() of anchors. - """ - on_pairs: set[tuple] = set() - anchor_ids: set[int] = set() + ) -> tuple[set[tuple[int, int]], set[PlaceableAsset]]: + """Return On-linked identity pairs and anchors.""" + on_pairs: set[tuple[int, int]] = set() + anchors: set[PlaceableAsset] = set() for obj in positions: for rel in obj.get_relations(): if isinstance(rel, On) and rel.parent in positions: on_pairs.add((id(obj), id(rel.parent))) on_pairs.add((id(rel.parent), id(obj))) if obj.is_anchor: - anchor_ids.add(id(obj)) - return on_pairs, anchor_ids + anchors.add(obj) + return on_pairs, anchors def _non_skip_pairs( self, @@ -428,14 +458,14 @@ def _non_skip_pairs( skip_mesh_pairs: bool = False, ) -> Iterator[tuple[PlaceableAsset, PlaceableAsset]]: """Yield non-relation object pairs, optionally skipping pairs handled by mesh collision.""" - on_pairs, anchor_ids = self._collect_skip_pairs(positions) + on_pairs, anchors = self._collect_skip_pairs(positions) mesh_manager = self._get_cpu_mesh_manager() if skip_mesh_pairs else None default_collision_mode = self._params.solver_params.collision_mode objects = list(positions.keys()) for i in range(len(objects)): for j in range(i + 1, len(objects)): a, b = objects[i], objects[j] - if id(a) in anchor_ids and id(b) in anchor_ids: + if a in anchors and b in anchors: continue if (id(a), id(b)) in on_pairs: continue @@ -455,18 +485,22 @@ def _non_skip_pairs( def _validate_no_overlap( self, positions: dict[PlaceableAsset, tuple[float, float, float]], - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox], + rotations: dict[PlaceableAsset, tuple[float, float, float, float]] | None = None, collision_objects: list[CollisionObject] | None = None, skip_mesh_pairs: bool = False, ) -> bool: - """AABB overlap check on pre-rotated env_bboxes. Skips On-pairs and anchor-anchor pairs.""" + """OBB penetration check using the same geometry as the solver.""" 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) + rotations = rotations or {} + _, anchors = self._collect_skip_pairs(positions) for a, b in self._non_skip_pairs(positions, 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): + a_world = _candidate_world_bbox(a, env_bboxes[a], positions[a], rotations) + b_world = _candidate_world_bbox(b, env_bboxes[b], positions[b], rotations) + if a_world.penetration(b_world, clearance_m=margin).item() > 0.0: if self._params.verbose: print(f" Overlap between '{a.name}' and '{b.name}'") return False @@ -477,9 +511,9 @@ def _validate_no_overlap( mesh_manager = self._get_cpu_mesh_manager() if skip_mesh_pairs else None default_collision_mode = self._params.solver_params.collision_mode for obj in positions: - if id(obj) in anchor_ids: + if obj in anchors: continue - obj_world = env_bboxes[obj].translated(positions[obj]) + obj_world = _candidate_world_bbox(obj, env_bboxes[obj], positions[obj], rotations) for background, background_world in background_worlds: if ( mesh_manager is not None @@ -487,7 +521,7 @@ def _validate_no_overlap( and mesh_manager.get_collision_mesh(background) is not None ): continue - if obj_world.overlaps(background_world, margin=margin).item(): + if obj_world.penetration(background_world, clearance_m=margin).item() > 0.0: if self._params.verbose: print(f" Overlap between '{obj.name}' and background '{background.name}'") return False @@ -507,17 +541,17 @@ def _get_cpu_mesh_manager(self) -> WarpMeshAndSphereCache: def _validate_no_overlap_mesh( self, positions: dict[PlaceableAsset, tuple[float, float, float]], - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox], - orientations: dict[PlaceableAsset, float] | None = None, + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox], + rotations: dict[PlaceableAsset, tuple[float, float, float, float]] | None = None, collision_objects: list[CollisionObject] | None = None, ) -> bool: - """Sphere-to-SDF overlap check; both-meshless pairs fall back to AABB validation.""" + """Sphere-to-SDF overlap check; meshless pairs use OBB validation.""" 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 [] + rotations = rotations or {} default_collision_mode = self._params.solver_params.collision_mode for a, b in self._non_skip_pairs(positions): @@ -531,43 +565,43 @@ def _validate_no_overlap_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" + " falling back to OBB validation for this pair" ) continue - a_pos = torch.tensor(positions[a], dtype=torch.float32) - b_pos = torch.tensor(positions[b], dtype=torch.float32) + a_pos, a_rotation = self._candidate_mesh_pose(a, positions[a], rotations) + b_pos, b_rotation = self._candidate_mesh_pose(b, positions[b], rotations) + a_collision_mesh = self._collision_mesh_or_bbox_proxy(a_mesh, env_bboxes[a]) + b_collision_mesh = self._collision_mesh_or_bbox_proxy(b_mesh, env_bboxes[b]) - if b_mesh is not None and self._spheres_penetrate_mesh( + if self._spheres_penetrate_mesh( a, - self._collision_mesh_or_aabb_proxy(a_mesh, env_bboxes[a]), + a_collision_mesh, a if a_mesh is not None else None, - a_mesh is not None, - a.is_anchor, a_pos, + a_rotation, b, - b_mesh, + b_collision_mesh, + b if b_mesh is not None else None, b_pos, - b.is_anchor, + b_rotation, mesh_manager, tolerance, - orientations, ): return False - if a_mesh is not None and self._spheres_penetrate_mesh( + if self._spheres_penetrate_mesh( b, - self._collision_mesh_or_aabb_proxy(b_mesh, env_bboxes[b]), + b_collision_mesh, b if b_mesh is not None else None, - b_mesh is not None, - b.is_anchor, b_pos, + b_rotation, a, - a_mesh, + a_collision_mesh, + a if a_mesh is not None else None, a_pos, - a.is_anchor, + a_rotation, mesh_manager, tolerance, - orientations, ): return False @@ -579,7 +613,7 @@ def _validate_no_overlap_mesh( if object_uses_mesh_collision(source, default_collision_mode) else None ) - source_pos = torch.tensor(positions[source], dtype=torch.float32) + source_pos, source_rotation = self._candidate_mesh_pose(source, positions[source], rotations) for background in collision_objects: target_mesh = ( mesh_manager.get_collision_mesh(background) @@ -593,133 +627,88 @@ def _validate_no_overlap_mesh( target_pose, Pose ), f"Background collision object '{background.name}' must have a fixed Pose in MESH mode." target_pos = torch.tensor(target_pose.position_xyz, dtype=torch.float32) + target_rotation = torch.tensor(target_pose.rotation_xyzw, dtype=torch.float32) if self._spheres_penetrate_mesh( source, - self._collision_mesh_or_aabb_proxy(source_mesh, env_bboxes[source]), + self._collision_mesh_or_bbox_proxy(source_mesh, env_bboxes[source]), source if source_mesh is not None else None, - source_mesh is not None, - source.is_anchor, source_pos, + source_rotation, background, target_mesh, + background, target_pos, - True, + target_rotation, mesh_manager, tolerance, - orientations, ): return False return True @staticmethod - def _collision_mesh_or_aabb_proxy( + def _collision_mesh_or_bbox_proxy( mesh: trimesh.Trimesh | None, - bbox: AxisAlignedBoundingBox, + bbox: OrientedBoundingBox, ) -> trimesh.Trimesh: - """Return an object's collision mesh, or a box mesh matching the candidate AABB.""" + """Return a collision mesh or asset-local box proxy for the base OBB.""" if mesh is not None: return mesh - box_mesh = trimesh.creation.box(extents=bbox.size[0].detach().cpu().numpy()) - box_mesh.apply_translation(bbox.center[0].detach().cpu().numpy()) + box_mesh = trimesh.creation.box(extents=(2.0 * bbox.half_extents[0]).detach().cpu().numpy()) + rotation_xyzw = bbox.rotation_xyzw[0].detach().cpu().numpy() + transform = trimesh.transformations.quaternion_matrix(np.roll(rotation_xyzw, 1)) + transform[:3, 3] = bbox.center[0].detach().cpu().numpy() + box_mesh.apply_transform(transform) return box_mesh + @staticmethod + def _candidate_mesh_pose( + obj: PlaceableAsset, + position: tuple[float, float, float], + rotations: dict[PlaceableAsset, tuple[float, float, float, float]], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return the current world position and xyzw rotation used by mesh validation.""" + if obj.is_anchor: + pose = obj.get_initial_pose() + assert isinstance(pose, Pose), f"Anchor '{obj.name}' must have a fixed Pose in MESH mode." + return torch.tensor(pose.position_xyz, dtype=torch.float32), torch.tensor( + pose.rotation_xyzw, dtype=torch.float32 + ) + return torch.tensor(position, dtype=torch.float32), torch.tensor( + rotations.get(obj, (0.0, 0.0, 0.0, 1.0)), dtype=torch.float32 + ) + def _spheres_penetrate_mesh( self, source: PlaceableAsset, source_mesh: trimesh.Trimesh, source_sphere_cache_obj: PlaceableAsset | None, - source_applies_yaw: bool, - source_uses_pose_yaw: bool, source_pos: torch.Tensor, + source_rotation: torch.Tensor, target: PlaceableAsset | CollisionObject, target_mesh: trimesh.Trimesh, + target_mesh_cache_obj: PlaceableAsset | CollisionObject | None, target_pos: torch.Tensor, - target_uses_pose_yaw: bool, + target_rotation: torch.Tensor, mesh_manager: WarpMeshAndSphereCache, tolerance: float, - orientations: dict[PlaceableAsset, float] | None, ) -> bool: - """True if source's spheres penetrate target's mesh or if BVH returns no-face sentinel. - - source_applies_yaw describes whether sphere centers need sampled-yaw rotation. - *_uses_pose_yaw controls whether fixed anchors/passive obstacles contribute pose yaw. - """ + """Return whether source spheres penetrate the target mesh.""" spheres = mesh_manager.get_query_spheres(source_mesh, obj=source_sphere_cache_obj) - warp_mesh = mesh_manager.get_warp_mesh(target_mesh, obj=target) - centers = self._centers_in_target_frame( + warp_mesh = mesh_manager.get_warp_mesh(target_mesh, obj=target_mesh_cache_obj) + centers = transform_points_between_frames( spheres[:, :3], - source, - target, source_pos, + source_rotation, target_pos, - orientations, - source_applies_yaw=source_applies_yaw, - source_uses_pose_yaw=source_uses_pose_yaw, - target_uses_pose_yaw=target_uses_pose_yaw, + target_rotation, ) sdf = mesh_sdf(centers, warp_mesh) - mesh_manager.warn_sdf_sentinel(sdf) - if has_sdf_sentinel(sdf): - return True + assert not has_sdf_sentinel( + sdf + ), "MESH collision query could not resolve a target face; the promised collision geometry is unsupported." if (sdf < spheres[:, 3] + tolerance).any(): if self._params.verbose: print(f" Mesh overlap between '{source.name}' and '{target.name}'") return True return False - - @staticmethod - def _effective_yaw( - obj: PlaceableAsset | CollisionObject, - orientations: dict[PlaceableAsset, float] | None, - use_pose_yaw: bool, - ) -> float: - """Resolve effective Z-yaw from sampled orientations or, when allowed, fixed initial pose.""" - if orientations is not None and obj in orientations: - return orientations[cast("PlaceableAsset", obj)] - if not use_pose_yaw: - return 0.0 - pose = obj.get_initial_pose() - if not isinstance(pose, Pose): - return 0.0 - return yaw_from_quat_xyzw(pose.rotation_xyzw) - - @staticmethod - def _pair_aabb_overlaps( - a_bbox: AxisAlignedBoundingBox, - b_bbox: AxisAlignedBoundingBox, - pos_a: tuple[float, float, float], - pos_b: tuple[float, float, float], - yaw_a: float, - yaw_b: float, - margin: float, - ) -> bool: - """True if two yaw-rotated, world-translated AABBs overlap (with margin).""" - if yaw_a != 0.0: - a_bbox = a_bbox.rotated_around_z(yaw_a) - if yaw_b != 0.0: - b_bbox = b_bbox.rotated_around_z(yaw_b) - a_world = a_bbox.translated(pos_a) - b_world = b_bbox.translated(pos_b) - return a_world.overlaps(b_world, margin=margin).item() - - @staticmethod - def _centers_in_target_frame( - centers_local: torch.Tensor, - source_obj: PlaceableAsset, - target_obj: PlaceableAsset | CollisionObject, - source_pos: torch.Tensor, - target_pos: torch.Tensor, - orientations: dict[PlaceableAsset, float] | None, - source_applies_yaw: bool = True, - source_uses_pose_yaw: bool = True, - target_uses_pose_yaw: bool = True, - ) -> torch.Tensor: - """Transform source sphere centers into the target's local frame (Z-yaw only).""" - src_yaw = ( - NoOverlapValidator._effective_yaw(source_obj, orientations, source_uses_pose_yaw) - if source_applies_yaw - else 0.0 - ) - tgt_yaw = NoOverlapValidator._effective_yaw(target_obj, orientations, target_uses_pose_yaw) - return centers_in_target_frame(centers_local, src_yaw, tgt_yaw, source_pos - target_pos) diff --git a/isaaclab_arena/relations/relation_loss_strategies.py b/isaaclab_arena/relations/relation_loss_strategies.py index 720d399b0f..8ae28bfda1 100644 --- a/isaaclab_arena/relations/relation_loss_strategies.py +++ b/isaaclab_arena/relations/relation_loss_strategies.py @@ -12,13 +12,12 @@ 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, ) from isaaclab_arena.relations.relations import Side -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox if TYPE_CHECKING: from isaaclab_arena.relations.relations import ( @@ -49,16 +48,13 @@ class Direction(IntEnum): @dataclass(frozen=True) class SideConfig: - """Configuration for computing NextTo loss for a given axis direction. - - Attributes: - primary_axis: Axis along which child is placed (X or Y). - direction: POSITIVE if child should be in positive direction from parent, - NEGATIVE if child should be in negative direction. - """ + """Axis and direction of a NextTo side.""" primary_axis: Axis + """Axis along which the child is placed.""" + direction: Direction + """Direction from the parent.""" @property def band_axis(self) -> Axis: @@ -77,8 +73,8 @@ def band_axis(self) -> Axis: def next_to_violations( cfg: SideConfig, child_pos: torch.Tensor, - child_bbox: AxisAlignedBoundingBox, - parent_world_bbox: AxisAlignedBoundingBox, + child_bbox: OrientedBoundingBox, + parent_world_bbox: OrientedBoundingBox, distance_m: float, ) -> tuple[torch.Tensor, torch.Tensor]: """Side and distance violation magnitudes (meters, >= 0) for a NextTo relation. @@ -96,13 +92,17 @@ def next_to_violations( Returns: (half_plane, distance) tensors of shape (N,), each >= 0 and zero when satisfied. """ + axis = child_pos.new_zeros(3) + axis[cfg.primary_axis] = 1.0 + child_min, child_max = child_bbox.get_bounds_along_axis(axis) + parent_min, parent_max = parent_world_bbox.get_bounds_along_axis(axis) if cfg.direction == Direction.POSITIVE: - parent_edge = parent_world_bbox.max_point[:, cfg.primary_axis] - child_offset = child_bbox.min_point[:, cfg.primary_axis] + parent_edge = parent_max + child_offset = child_min penalty_side = "less" else: - parent_edge = parent_world_bbox.min_point[:, cfg.primary_axis] - child_offset = child_bbox.max_point[:, cfg.primary_axis] + parent_edge = parent_min + child_offset = child_max penalty_side = "greater" primary = child_pos[:, cfg.primary_axis] @@ -115,8 +115,8 @@ def next_to_violations( def not_next_to_violations( cfg: SideConfig, child_pos: torch.Tensor, - child_bbox: AxisAlignedBoundingBox, - parent_world_bbox: AxisAlignedBoundingBox, + child_bbox: OrientedBoundingBox, + parent_world_bbox: OrientedBoundingBox, margin_m: float, ) -> tuple[torch.Tensor, torch.Tensor]: """Per-route escape distances (meters, >= 0) for a NotNextTo relation. @@ -135,17 +135,22 @@ def not_next_to_violations( Returns: (remaining_side, remaining_cross) tensors of shape (N,), each >= 0. """ + primary_axis = child_pos.new_zeros(3) + primary_axis[cfg.primary_axis] = 1.0 + parent_primary_min, parent_primary_max = parent_world_bbox.get_bounds_along_axis(primary_axis) if cfg.direction == Direction.POSITIVE: - parent_edge = parent_world_bbox.max_point[:, cfg.primary_axis] + parent_edge = parent_primary_max blocked_side_penalty = "greater" else: - parent_edge = parent_world_bbox.min_point[:, cfg.primary_axis] + parent_edge = parent_primary_min blocked_side_penalty = "less" - parent_band_min = parent_world_bbox.min_point[:, cfg.band_axis] - parent_band_max = parent_world_bbox.max_point[:, cfg.band_axis] - valid_band_min = parent_band_min - child_bbox.min_point[:, cfg.band_axis] - valid_band_max = parent_band_max - child_bbox.max_point[:, cfg.band_axis] + band_axis = child_pos.new_zeros(3) + band_axis[cfg.band_axis] = 1.0 + child_band_min, child_band_max = child_bbox.get_bounds_along_axis(band_axis) + parent_band_min, parent_band_max = parent_world_bbox.get_bounds_along_axis(band_axis) + valid_band_min = parent_band_min - child_band_min + valid_band_max = parent_band_max - child_band_max primary = child_pos[:, cfg.primary_axis] cross = child_pos[:, cfg.band_axis] @@ -169,14 +174,13 @@ def compute_loss( self, relation: UnaryRelation, child_pos: torch.Tensor, - child_bbox: AxisAlignedBoundingBox, + child_bbox: OrientedBoundingBox, ) -> torch.Tensor: """Compute the loss for a unary relation constraint. Args: relation: The relation object containing constraint metadata. - child_pos: Child object position tensor. Accepts (3,) for single-env - backward compat or (N, 3) for batched. + child_pos: Child position with shape (3,) or (N, 3). child_bbox: Child object local bounding box (N=1). Returns: @@ -193,15 +197,14 @@ def compute_loss( self, relation: Relation, child_pos: torch.Tensor, - child_bbox: AxisAlignedBoundingBox, - parent_world_bbox: AxisAlignedBoundingBox, + child_bbox: OrientedBoundingBox, + parent_world_bbox: OrientedBoundingBox, ) -> torch.Tensor: """Compute the loss for a relation constraint. Args: relation: The relation object containing relationship metadata. - child_pos: Child object position tensor. Accepts (3,) for single-env - backward compat or (N, 3) for batched. + child_pos: Child position with shape (3,) or (N, 3). child_bbox: Child object local bounding box (N=1). parent_world_bbox: Parent bounding box in world coordinates. @@ -234,8 +237,8 @@ def compute_loss( self, relation: NextTo, child_pos: torch.Tensor, - child_bbox: AxisAlignedBoundingBox, - parent_world_bbox: AxisAlignedBoundingBox, + child_bbox: OrientedBoundingBox, + parent_world_bbox: OrientedBoundingBox, ) -> torch.Tensor: """Compute loss for NextTo relation. @@ -264,10 +267,12 @@ def compute_loss( distance_loss = self.slope * distance_raw # 2. Band position loss: child placed at target position within parent's perpendicular extent - parent_band_min = parent_world_bbox.min_point[:, cfg.band_axis] - parent_band_max = parent_world_bbox.max_point[:, cfg.band_axis] - valid_band_min = parent_band_min - child_bbox.min_point[:, cfg.band_axis] - valid_band_max = parent_band_max - child_bbox.max_point[:, cfg.band_axis] + band_axis = child_pos.new_zeros(3) + band_axis[cfg.band_axis] = 1.0 + child_band_min, child_band_max = child_bbox.get_bounds_along_axis(band_axis) + parent_band_min, parent_band_max = parent_world_bbox.get_bounds_along_axis(band_axis) + valid_band_min = parent_band_min - child_band_min + valid_band_max = parent_band_max - child_band_max # Convert cross_position_ratio [-1, 1] to interpolation factor [0, 1]: -1 = min, 0 = center, 1 = max t = (relation.cross_position_ratio + 1.0) / 2.0 target_band_pos = valid_band_min + t * (valid_band_max - valid_band_min) @@ -315,8 +320,8 @@ def compute_loss( self, relation: On, child_pos: torch.Tensor, - child_bbox: AxisAlignedBoundingBox, - parent_world_bbox: AxisAlignedBoundingBox, + child_bbox: OrientedBoundingBox, + parent_world_bbox: OrientedBoundingBox, ) -> torch.Tensor: """Compute loss for On relation. @@ -333,20 +338,21 @@ def compute_loss( if single_input: child_pos = child_pos.unsqueeze(0) - # Parent world-space extents from the world bounding box - parent_x_min = parent_world_bbox.min_point[:, 0] - parent_x_max = parent_world_bbox.max_point[:, 0] - parent_y_min = parent_world_bbox.min_point[:, 1] - parent_y_max = parent_world_bbox.max_point[:, 1] - parent_z_max = parent_world_bbox.max_point[:, 2] # Top surface + axes = torch.eye(3, dtype=child_pos.dtype, device=child_pos.device) + parent_x_min, parent_x_max = parent_world_bbox.get_bounds_along_axis(axes[0]) + parent_y_min, parent_y_max = parent_world_bbox.get_bounds_along_axis(axes[1]) + _, parent_z_max = parent_world_bbox.get_bounds_along_axis(axes[2]) + child_x_min, child_x_max = child_bbox.get_bounds_along_axis(axes[0]) + child_y_min, child_y_max = child_bbox.get_bounds_along_axis(axes[1]) + child_z_min, _ = child_bbox.get_bounds_along_axis(axes[2]) # Compute valid position ranges such that child's entire footprint is within parent, # with the parent's extent inset by edge_margin_m so the footprint stays off the rim. m = relation.edge_margin_m - valid_x_min = parent_x_min + m - child_bbox.min_point[:, 0] # child's left at parent's left + margin - valid_x_max = parent_x_max - m - child_bbox.max_point[:, 0] # child's right at parent's right - margin - valid_y_min = parent_y_min + m - child_bbox.min_point[:, 1] - valid_y_max = parent_y_max - m - child_bbox.max_point[:, 1] + valid_x_min = parent_x_min + m - child_x_min + valid_x_max = parent_x_max - m - child_x_max + valid_y_min = parent_y_min + m - child_y_min + valid_y_max = parent_y_max - m - child_y_max # The bounds invert (lower > upper) when the margin is too large for the surface or the # child is oversized. The loss becomes a non-zero constant with gradient zero. @@ -368,7 +374,7 @@ def compute_loss( ) # 3. Z point loss: child bottom = parent top + clearance - target_z = parent_z_max + relation.clearance_m - child_bbox.min_point[:, 2] + target_z = parent_z_max + relation.clearance_m - child_z_min z_loss = single_point_linear_loss(child_pos[:, 2], target_z, slope=self.slope) if self.debug and child_pos.shape[0] == 1: @@ -424,8 +430,8 @@ def compute_loss( self, relation: NotNextTo, child_pos: torch.Tensor, - child_bbox: AxisAlignedBoundingBox, - parent_world_bbox: AxisAlignedBoundingBox, + child_bbox: OrientedBoundingBox, + parent_world_bbox: OrientedBoundingBox, ) -> torch.Tensor: """Compute loss for ``NotNextTo``.""" single_input = child_pos.dim() == 1 @@ -457,51 +463,27 @@ def compute_loss( class NoCollisionLossStrategy: - """AABB no-overlap loss between object pairs.""" + """Scale OBB minimum-penetration distances into collision loss.""" + + def __init__(self, slope: float = 10.0): + """Initialize the overlap-loss scale. - def __init__( - self, - slope: float = 10.0, - ): - """ Args: slope: Gradient magnitude for overlap loss. """ self.slope = slope - def compute_loss_batched( - self, - clearance_m: float, - subject_min: torch.Tensor, - subject_max: torch.Tensor, - obstacle_min: torch.Tensor, - obstacle_max: torch.Tensor, - ) -> torch.Tensor: - """Overlap-volume no-overlap loss for boxes already reduced to world-space extents. + def compute_loss_batched(self, penetration: torch.Tensor) -> torch.Tensor: + """Return scaled minimum penetration for directed pairs. Args: - clearance_m: Minimum clearance between boxes in meters. - subject_min: World-space min extent of the subject box, shape (num_pairs, batch_size, 3). - subject_max: World-space max extent of the subject box, shape (num_pairs, batch_size, 3). - obstacle_min: World-space min extent of the obstacle box, shape (num_pairs, batch_size, 3). - obstacle_max: World-space max extent of the obstacle box, shape (num_pairs, batch_size, 3). + penetration: OBB penetration distances with shape (P, N). Returns: - Per-pair, per-env loss of shape (num_pairs, batch_size). + Per-pair, per-env loss with shape (P, N). """ - 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) + assert penetration.ndim == 2, f"Expected penetration shape (P, N), got {tuple(penetration.shape)}." + return self.slope * penetration class AtPositionLossStrategy(UnaryRelationLossStrategy): @@ -523,7 +505,7 @@ def compute_loss( self, relation: AtPosition, child_pos: torch.Tensor, - child_bbox: AxisAlignedBoundingBox, + child_bbox: OrientedBoundingBox, ) -> torch.Tensor: """Compute loss for AtPosition relation. @@ -576,7 +558,7 @@ def compute_loss( self, relation: PositionLimits, child_pos: torch.Tensor, - child_bbox: AxisAlignedBoundingBox, + child_bbox: OrientedBoundingBox, ) -> torch.Tensor: """Compute loss for PositionLimits relation. diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py index de702a41b5..7e83caedfc 100644 --- a/isaaclab_arena/relations/relation_solver.py +++ b/isaaclab_arena/relations/relation_solver.py @@ -10,8 +10,8 @@ from typing import TYPE_CHECKING, cast from isaaclab_arena.relations.collision_mode import CollisionMode, get_object_collision_mode -from isaaclab_arena.relations.no_overlap_aabb import compute_no_overlap_loss_aabb from isaaclab_arena.relations.no_overlap_mesh import compute_no_overlap_loss_mesh, prepare_mesh_collision_cache +from isaaclab_arena.relations.no_overlap_obb import compute_no_overlap_loss_obb from isaaclab_arena.relations.relation_loss_strategies import ( NoCollisionLossStrategy, RelationLossStrategy, @@ -26,7 +26,7 @@ from isaaclab_arena.relations.mesh_pair_cache import MeshPairCache 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 + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox class RelationSolver: @@ -48,13 +48,11 @@ 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(slope=self.params.collision_loss_slope) self._last_loss_history: list[float] = [] self._last_position_history: list = [] self._last_loss_per_env: torch.Tensor | None = None self._last_no_overlap_pair_count: int = 0 - self._mesh_orientations: list[dict[PlaceableAsset, float]] | None = None self._warned_no_mesh: set[str] = set() self._mesh_manager: WarpMeshAndSphereCache | None = None self._mesh_cache: MeshPairCache | None = None @@ -154,12 +152,11 @@ def _compute_no_overlap_loss( state, self._mesh_cache, self._mesh_manager, - self._mesh_orientations, - self.params.clearance_m, - self._no_collision_strategy.slope, - debug, + clearance_m=self.params.clearance_m, + slope=self._no_collision_strategy.slope, + debug=debug, ) - aabb_loss, n = compute_no_overlap_loss_aabb( + bbox_loss, pair_count = compute_no_overlap_loss_obb( state, self._no_collision_strategy, self.params.clearance_m, @@ -168,9 +165,9 @@ def _compute_no_overlap_loss( skip_mesh_pairs=True, debug=debug, ) - self._last_no_overlap_pair_count = n - return mesh_loss + aabb_loss - loss, n = compute_no_overlap_loss_aabb( + self._last_no_overlap_pair_count = pair_count + return mesh_loss + bbox_loss + loss, pair_count = compute_no_overlap_loss_obb( state, self._no_collision_strategy, self.params.clearance_m, @@ -178,16 +175,15 @@ def _compute_no_overlap_loss( self.params.collision_mode, debug=debug, ) - self._last_no_overlap_pair_count = n + self._last_no_overlap_pair_count = pair_count return loss def solve( self, objects: list[PlaceableAsset], initial_positions: list[dict[PlaceableAsset, tuple[float, float, float]]], - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox] | None = None, - env_bboxes_include_yaw: bool = False, - orientations: list[dict[PlaceableAsset, float]] | None = None, + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox] | None = None, + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]] | None = None, collision_objects: list[CollisionObject] | None = None, ) -> list[dict[PlaceableAsset, tuple[float, float, float]]]: """Solve for optimal positions of all objects. @@ -199,13 +195,10 @@ def solve( for single-env placement. env_bboxes: Optional per-env bounding boxes keyed by object. ObjectPlacer always supplies these, with each - AxisAlignedBoundingBox shaped (batch, 3). Direct solver calls + OrientedBoundingBox shaped with N=batch. Direct solver calls may omit them to use each object's default get_bounding_box(). - env_bboxes_include_yaw: Whether env_bboxes already enclose each object's - yawed footprint. ObjectPlacer sets this after applying candidate yaw; - direct callers should leave it False unless they expanded the bboxes. - orientations: Optional per-env yaw angles (radians about Z) per object. - Used in MESH mode to rotate sphere centers before collision queries. + rotations: Optional fixed per-candidate xyzw rotations for movable + objects. FaceTo rotations are derived from current positions. collision_objects: Optional fixed background obstacles included in the no-overlap collision term only. They are not optimized and carry no relation constraints. @@ -213,10 +206,14 @@ def solve( Returns: List of dicts (one per env) mapping objects to their solved (x, y, z) positions. """ - assert not env_bboxes_include_yaw or env_bboxes is not None, "env_bboxes_include_yaw=True requires env_bboxes." device = torch.device("cuda" if torch.cuda.is_available() else "cpu") state = RelationSolverState( - objects, initial_positions, device=device, env_bboxes=env_bboxes, collision_objects=collision_objects + objects, + initial_positions, + device=device, + env_bboxes=env_bboxes, + rotations=rotations, + collision_objects=collision_objects, ) if self.params.verbose: @@ -241,7 +238,6 @@ def solve( torch.cuda.synchronize() solve_start = time.perf_counter() - # Precompute mesh collision cache (once per solve, before opt loop) self._mesh_collision_enabled = self._should_use_mesh_collision(state) if self._mesh_collision_enabled: non_anchor_objects = state.optimizable_objects @@ -252,7 +248,6 @@ def solve( if isinstance(rel, On): on_pairs.add((id(obj), id(rel.parent))) on_pairs.add((id(rel.parent), id(obj))) - self._mesh_orientations = orientations device_str = str(state.device) if self._mesh_manager is None or self._mesh_manager.device != device_str: from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache @@ -264,11 +259,8 @@ def solve( on_pairs, self._warned_no_mesh, default_collision_mode=self.params.collision_mode, - bboxes_include_yaw=env_bboxes_include_yaw, ) - self._mesh_manager.reset_sentinel_warning() else: - self._mesh_orientations = None self._mesh_cache = None # Setup optimizer (only for optimizable positions) @@ -387,39 +379,35 @@ def _print_relation_debug( """Print debug information for a single binary relation.""" child_bbox = obj.get_bounding_box() parent_world_bbox = relation.parent.get_world_bounding_box() + child_min, child_max = child_bbox.get_axis_aligned_bounds() + parent_min, parent_max = parent_world_bbox.get_axis_aligned_bounds() print(f"\n=== {obj.name} -> {type(relation).__name__}({relation.parent.name}) ===") print(f" Child pos: ({child_pos[0].item():.4f}, {child_pos[1].item():.4f}, {child_pos[2].item():.4f})") print( - f" Child bbox: min={child_bbox.min_point[0].tolist()}, max={child_bbox.max_point[0].tolist()}," - f" size={child_bbox.size[0].tolist()}" + f" Child bbox: min={child_min[0].tolist()}, max={child_max[0].tolist()}," + f" size={(2.0 * child_bbox.half_extents[0]).tolist()}" ) print(f" Parent pos: ({parent_pos[0].item():.4f}, {parent_pos[1].item():.4f}, {parent_pos[2].item():.4f})") print( - f" Parent world bbox: min={parent_world_bbox.min_point[0].tolist()}," - f" max={parent_world_bbox.max_point[0].tolist()}, size={parent_world_bbox.size[0].tolist()}" + f" Parent world bbox: min={parent_min[0].tolist()}," + f" max={parent_max[0].tolist()}, size={(2.0 * parent_world_bbox.half_extents[0]).tolist()}" ) # Child world extents child_x_range = ( - child_pos[0].item() + child_bbox.min_point[0, 0].item(), - child_pos[0].item() + child_bbox.max_point[0, 0].item(), + child_pos[0].item() + child_min[0, 0].item(), + child_pos[0].item() + child_max[0, 0].item(), ) child_y_range = ( - child_pos[1].item() + child_bbox.min_point[0, 1].item(), - child_pos[1].item() + child_bbox.max_point[0, 1].item(), + child_pos[1].item() + child_min[0, 1].item(), + child_pos[1].item() + child_max[0, 1].item(), ) print(f" Child world X: [{child_x_range[0]:.4f}, {child_x_range[1]:.4f}]") print(f" Child world Y: [{child_y_range[0]:.4f}, {child_y_range[1]:.4f}]") - print( - f" Parent world X: [{parent_world_bbox.min_point[0, 0].item():.4f}," - f" {parent_world_bbox.max_point[0, 0].item():.4f}]" - ) - print( - f" Parent world Y: [{parent_world_bbox.min_point[0, 1].item():.4f}," - f" {parent_world_bbox.max_point[0, 1].item():.4f}]" - ) + print(f" Parent world X: [{parent_min[0, 0].item():.4f}, {parent_max[0, 0].item():.4f}]") + print(f" Parent world Y: [{parent_min[0, 1].item():.4f}, {parent_max[0, 1].item():.4f}]") print(f" Loss: {loss.item():.6f}") @@ -431,13 +419,14 @@ def _print_unary_relation_debug( ) -> None: """Print debug information for a unary relation (no parent).""" child_bbox = obj.get_bounding_box() + child_min, child_max = child_bbox.get_axis_aligned_bounds() params = {k: v for k, v in relation.__dict__.items() if v is not None and k != "relation_loss_weight"} param_str = ", ".join(f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}" for k, v in params.items()) print(f"\n=== {obj.name} -> {type(relation).__name__}({param_str}) ===") print(f" Child pos: ({child_pos[0].item():.4f}, {child_pos[1].item():.4f}, {child_pos[2].item():.4f})") print( - f" Child bbox: min={child_bbox.min_point[0].tolist()}, max={child_bbox.max_point[0].tolist()}," - f" size={child_bbox.size[0].tolist()}" + f" Child bbox: min={child_min[0].tolist()}, max={child_max[0].tolist()}," + f" size={(2.0 * child_bbox.half_extents[0]).tolist()}" ) print(f" Loss: {loss.item():.6f}") diff --git a/isaaclab_arena/relations/relation_solver_params.py b/isaaclab_arena/relations/relation_solver_params.py index 26b125663a..2df93dbfd3 100644 --- a/isaaclab_arena/relations/relation_solver_params.py +++ b/isaaclab_arena/relations/relation_solver_params.py @@ -62,6 +62,9 @@ class RelationSolverParams: The solver adds a no-overlap loss for all pairs automatically. Set to 0.0 to only reject actual overlaps (no safety margin).""" + collision_loss_slope: float = 1000.0 + """Scale applied to OBB and mesh penetration, ten times the strongest default relation slope.""" + # default_factory ensures each instance gets its own dict (mutable defaults are shared across instances) strategies: dict[type[RelationBase], RelationLossStrategy | UnaryRelationLossStrategy] = field( default_factory=_default_strategies @@ -70,3 +73,4 @@ class RelationSolverParams: def __post_init__(self): assert self.clearance_m >= 0, f"clearance_m must be >= 0, got {self.clearance_m}" + assert self.collision_loss_slope > 0, f"collision_loss_slope must be > 0, got {self.collision_loss_slope}" diff --git a/isaaclab_arena/relations/relation_solver_state.py b/isaaclab_arena/relations/relation_solver_state.py index 1d74242c48..19c337b5dd 100644 --- a/isaaclab_arena/relations/relation_solver_state.py +++ b/isaaclab_arena/relations/relation_solver_state.py @@ -8,8 +8,10 @@ import torch from typing import TYPE_CHECKING -from isaaclab_arena.relations.relations import get_anchor_objects -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.relations.relations import FaceTo, get_anchor_objects, get_relation +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox +from isaaclab_arena.utils.pose import Pose +from isaaclab_arena.utils.yaw import MINIMUM_FACING_DIRECTION_XY_M if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject @@ -17,21 +19,15 @@ class RelationSolverState: - """Encapsulates position state during optimization. - - This class manages the mapping between objects and their positions, - keeping anchor (fixed) and optimizable positions separate internally - while providing an interface for position lookups. - - Positions are always stored as (batch_size, num_objects, 3). - """ + """Batched object poses and bounding boxes used by the relation solver.""" def __init__( self, objects: list[PlaceableAsset], initial_positions: list[dict[PlaceableAsset, tuple[float, float, float]]], device: torch.device | None = None, - env_bboxes: dict[PlaceableAsset, AxisAlignedBoundingBox] | None = None, + env_bboxes: dict[PlaceableAsset, OrientedBoundingBox] | None = None, + rotations: list[dict[PlaceableAsset, tuple[float, float, float, float]]] | None = None, collision_objects: list[CollisionObject] | None = None, ): """Initialize optimization state. @@ -43,6 +39,8 @@ def __init__( length > 1 = batched. device: Torch device for all tensors. Defaults to CPU. env_bboxes: Optional per-env bounding boxes keyed by object. + rotations: Fixed candidate rotations in xyzw order. FaceTo subjects + derive their effective world-Z rotation from current subject/target positions. collision_objects: Optional fixed background obstacles that participate in no-overlap collision only (never in relation constraints). They keep a constant world bounding box and are not optimized. Must be disjoint from objects. @@ -60,40 +58,35 @@ def __init__( "both optimized and a fixed collision obstacle." ) - # Build object-to-index mapping self._obj_to_idx: dict[PlaceableAsset, int] = {obj: i for i, obj in enumerate(objects)} self._device = device or torch.device("cpu") self._batch_size = len(initial_positions) - # Validate that every dict contains all objects before building the tensor. for d in initial_positions: for obj in objects: assert obj in d, f"Missing initial position for {obj.name}" - # Build all positions as a single (N, num_objects, 3) tensor in one call. pos_nested = [[d[obj] for obj in objects] for d in initial_positions] all_positions = torch.tensor(pos_nested, dtype=torch.float32, device=self._device) - # Separate anchor positions from optimizable positions self._anchor_indices: set[int] = {self._obj_to_idx[obj] for obj in self._anchor_objects} - # Anchors must be identical across envs (they are fixed reference points). for idx in self._anchor_indices: - for env_idx in range(1, self._batch_size): - assert torch.allclose(all_positions[0, idx], all_positions[env_idx, idx]), ( - f"Anchor '{objects[idx].name}' has different positions across envs " - f"(env 0: {all_positions[0, idx].tolist()}, env {env_idx}: {all_positions[env_idx, idx].tolist()})" - ) + pose = objects[idx].get_initial_pose() + assert isinstance(pose, Pose), f"Anchor '{objects[idx].name}' must have a fixed Pose." + fixed_position = torch.tensor(pose.position_xyz, dtype=torch.float32, device=self._device) + assert torch.allclose(all_positions[:, idx], fixed_position.expand(self._batch_size, 3)), ( + f"Anchor '{objects[idx].name}' supplied positions must match its fixed Pose position " + f"{pose.position_xyz}, got {all_positions[:, idx].tolist()}." + ) self._anchor_positions: dict[int, torch.Tensor] = { idx: all_positions[0, idx].clone() for idx in self._anchor_indices } - # Pre-build anchor positions as (1, num_objects, 3) for fast _reconstruct_all_positions. self._anchor_pos_tensor = torch.zeros(1, len(objects), 3, dtype=torch.float32, device=self._device) for idx, pos in self._anchor_positions.items(): self._anchor_pos_tensor[0, idx, :] = pos - # Build optimizable positions tensor by slicing from the full tensor. self._optimizable_indices = [i for i in range(len(objects)) if i not in self._anchor_indices] self._global_to_opt_idx: dict[int, int] = { global_idx: opt_idx for opt_idx, global_idx in enumerate(self._optimizable_indices) @@ -107,13 +100,50 @@ def __init__( self._optimizable_positions = None self._env_bboxes = env_bboxes - - # Anchors and background collision objects are fixed, so their world bounding boxes are - # constant during the solve. Cache them once instead of recomputing every gradient step. - self._fixed_obstacle_world_bboxes: dict[PlaceableAsset | CollisionObject, AxisAlignedBoundingBox] = { - obj: obj.get_world_bounding_box().to(self._device) - for obj in (*self._anchor_objects, *self._collision_objects) - } + rotations = rotations or [{} for _ in range(self._batch_size)] + assert ( + len(rotations) == self._batch_size + ), f"rotations must contain one dictionary per candidate, got {len(rotations)} for N={self._batch_size}." + identity = (0.0, 0.0, 0.0, 1.0) + supplied_rotation_objects = {obj for candidate in rotations for obj in candidate} + invalid_rotation_objects = supplied_rotation_objects - set(self._optimizable_objects) + assert not invalid_rotation_objects, ( + "Rotation keys must belong to optimizable objects, got " + f"{sorted(obj.name for obj in invalid_rotation_objects)}." + ) + candidate_rotations: dict[PlaceableAsset, torch.Tensor] = {} + for obj in self._optimizable_objects: + values = [ + torch.as_tensor(candidate.get(obj, identity), dtype=torch.float32, device=self._device) + for candidate in rotations + ] + assert all( + value.shape == (4,) for value in values + ), f"Candidate rotations for '{obj.name}' must each have shape (4,)." + rotation_tensor = torch.stack(values) + assert rotation_tensor.shape == ( + self._batch_size, + 4, + ), f"Candidate rotations for '{obj.name}' must have shape (N, 4)." + assert torch.isfinite(rotation_tensor).all(), f"Candidate rotations for '{obj.name}' must be finite." + norms = torch.linalg.vector_norm(rotation_tensor, dim=-1) + assert torch.allclose( + norms, torch.ones_like(norms), atol=1e-5, rtol=1e-5 + ), f"Candidate rotations for '{obj.name}' must be unit quaternions." + candidate_rotations[obj] = rotation_tensor + self._base_rotations = {obj: candidate_rotations[obj] for obj in self._optimizable_objects} + + # Fixed world boxes do not depend on optimized positions. + self._fixed_obstacle_world_bboxes: dict[PlaceableAsset | CollisionObject, OrientedBoundingBox] = {} + for obj in self._anchor_objects: + pose = obj.get_initial_pose() + assert isinstance(pose, Pose), f"Anchor '{obj.name}' must have a fixed Pose." + local_bbox = env_bboxes[obj] if env_bboxes is not None and obj in env_bboxes else obj.get_bounding_box() + self._fixed_obstacle_world_bboxes[obj] = local_bbox.to(self._device).transformed( + pose.position_xyz, pose.rotation_xyzw + ) + for obj in self._collision_objects: + self._fixed_obstacle_world_bboxes[obj] = obj.get_world_bounding_box().to(self._device) @property def device(self) -> torch.device: @@ -156,32 +186,65 @@ def get_position(self, obj: PlaceableAsset) -> torch.Tensor: Returns: Position tensor of shape (batch_size, 3). - - Raises: - KeyError: If object is not tracked by this state. - RuntimeError: If requesting position for optimizable object when none exist. """ idx = self._obj_to_idx[obj] if idx in self._anchor_indices: return self._anchor_positions[idx].unsqueeze(0).expand(self._batch_size, 3) - if self._optimizable_positions is None: - raise RuntimeError(f"No optimizable positions available for object '{obj.name}'") + assert self._optimizable_positions is not None, f"No optimizable position for '{obj.name}'." opt_idx = self._global_to_opt_idx[idx] return self._optimizable_positions[:, opt_idx, :] - def get_fixed_obstacle_world_bbox(self, obj: PlaceableAsset | CollisionObject) -> AxisAlignedBoundingBox: + def get_fixed_obstacle_world_bbox(self, obj: PlaceableAsset | CollisionObject) -> OrientedBoundingBox: """Return the cached constant world bounding box for an anchor or collision object.""" assert ( obj in self._fixed_obstacle_world_bboxes ), f"'{obj.name}' is not a fixed obstacle (anchor or collision object) tracked by this state." return self._fixed_obstacle_world_bboxes[obj] - def get_bbox(self, obj: PlaceableAsset) -> AxisAlignedBoundingBox: - """Return the local bounding box for obj, moved to the state's device.""" + def get_rotation(self, obj: PlaceableAsset) -> torch.Tensor: + """Return effective candidate rotations with shape (N, 4), in xyzw order. + + Fixed candidate rotations come from RotateAroundSolution/random_yaw_init. + FaceTo is the one dynamic case: its world-Z quaternion is recomputed from + current positions on every loss evaluation, while positions remain the + solver's only optimized parameters. This position-derived rotation remains + differentiable, so geometry losses can propagate through it to both the + FaceTo subject and target positions. + """ + base = self._base_rotations[obj] + face_to = get_relation(obj, FaceTo) + if face_to is None: + return base + delta = self.get_position(face_to.parent)[:, :2] - self.get_position(obj)[:, :2] + distance = torch.linalg.vector_norm(delta, dim=1) + valid = distance > MINIMUM_FACING_DIRECTION_XY_M + safe_delta = torch.where(valid.unsqueeze(1), delta, delta.new_tensor([1.0, 0.0])) + yaw = torch.atan2(safe_delta[:, 1], safe_delta[:, 0]) + half_yaw = yaw * 0.5 + facing = torch.stack( + [ + torch.zeros_like(yaw), + torch.zeros_like(yaw), + torch.sin(half_yaw), + torch.cos(half_yaw), + ], + dim=1, + ) + return torch.where(valid.unsqueeze(1), facing, base) + + def get_base_bbox(self, obj: PlaceableAsset) -> OrientedBoundingBox: + """Return the asset-local bounding box for an object.""" if self._env_bboxes is not None and obj in self._env_bboxes: return self._env_bboxes[obj].to(self._device) return obj.get_bounding_box().to(self._device) + def get_bbox(self, obj: PlaceableAsset) -> OrientedBoundingBox: + """Return the candidate-oriented local bounding box for an object.""" + bbox = self.get_base_bbox(obj) + if obj in self._anchor_objects: + return bbox + return bbox.rotated_by_quat_unchecked(self.get_rotation(obj)) + def get_all_positions_snapshot(self) -> list[tuple[float, float, float]]: """Get detached copy of all positions for history tracking. diff --git a/isaaclab_arena/relations/warp_mesh_manager.py b/isaaclab_arena/relations/warp_mesh_manager.py index b8a16dc590..920c10f767 100644 --- a/isaaclab_arena/relations/warp_mesh_manager.py +++ b/isaaclab_arena/relations/warp_mesh_manager.py @@ -16,8 +16,6 @@ import warp as wp -from isaaclab_arena.relations.warp_sdf_kernels import has_sdf_sentinel, sdf_sentinel_count - if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject @@ -124,25 +122,8 @@ def __init__( self._warp_mesh_cache: dict[tuple, wp.Mesh] = {} self._sphere_cache: dict[tuple, torch.Tensor] = {} self._trimesh_cache: dict[tuple, trimesh.Trimesh | None] = {} - self._sentinel_warned: bool = False self._raw_open_mesh_warned: set[tuple] = set() - def reset_sentinel_warning(self) -> None: - """Re-arm for a new solve/validation pass.""" - self._sentinel_warned = False - - def warn_sdf_sentinel(self, sdf_values: torch.Tensor) -> None: - """Warn (once per pass) if any query hit the no-face sentinel.""" - if self._sentinel_warned: - return - if has_sdf_sentinel(sdf_values): - self._sentinel_warned = True - n_bad = sdf_sentinel_count(sdf_values) - print( - f" [MeshSDF] WARNING: {n_bad}/{len(sdf_values)} sphere queries returned sentinel SDF " - "(no mesh face found). Collision detection may be incomplete for these points." - ) - def get_collision_mesh(self, obj: CollisionObject) -> trimesh.Trimesh | None: """Return the cached collision mesh, extracting from USD on first access.""" from isaaclab_arena.assets.object import Object diff --git a/isaaclab_arena/relations/warp_sdf_kernels.py b/isaaclab_arena/relations/warp_sdf_kernels.py index 670ffb91c9..bfc8423a79 100644 --- a/isaaclab_arena/relations/warp_sdf_kernels.py +++ b/isaaclab_arena/relations/warp_sdf_kernels.py @@ -113,16 +113,6 @@ def has_sdf_sentinel(sdf_values: torch.Tensor) -> bool: return bool((sdf_values >= _SDF_SENTINEL).any()) -def sdf_sentinel_count(sdf_values: torch.Tensor) -> int: - """Number of queries that hit the no-face sentinel.""" - return int((sdf_values >= _SDF_SENTINEL).sum().item()) - - -def clamp_sdf_sentinel(sdf_values: torch.Tensor) -> torch.Tensor: - """Replace sentinel SDF values with 0 so no-face hits produce a constant penalty with zero gradient.""" - return torch.where(sdf_values >= _SDF_SENTINEL, torch.zeros_like(sdf_values), sdf_values) - - # --------------------------------------------------------------------------- # Multi-mesh kernel: query multiple meshes in a single launch # --------------------------------------------------------------------------- diff --git a/isaaclab_arena/tests/dummy_embodiment.py b/isaaclab_arena/tests/dummy_embodiment.py index 0c66136f11..0d376a7bef 100644 --- a/isaaclab_arena/tests/dummy_embodiment.py +++ b/isaaclab_arena/tests/dummy_embodiment.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING from isaaclab_arena.relations.placement_asset import PlaceableAsset -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose, PosePerEnv if TYPE_CHECKING: @@ -28,7 +28,7 @@ class DummyEmbodiment(PlaceableAsset): def __init__( self, name: str, - bounding_box: AxisAlignedBoundingBox, + bounding_box: OrientedBoundingBox, initial_pose: Pose | None = None, collision_mesh: trimesh.Trimesh | None = None, scene_name: str | None = None, @@ -58,7 +58,7 @@ def _build_reset_event(self) -> EventTermCfg | None: params={"scene_writes": self.layout_pose_to_scene_writes(self.initial_pose)}, ) - def get_bounding_box(self) -> AxisAlignedBoundingBox: + def get_bounding_box(self) -> OrientedBoundingBox: """Return root-relative bounds.""" return self.bounding_box diff --git a/isaaclab_arena/tests/dummy_object.py b/isaaclab_arena/tests/dummy_object.py index 77e164c13c..ee646a352a 100644 --- a/isaaclab_arena/tests/dummy_object.py +++ b/isaaclab_arena/tests/dummy_object.py @@ -10,7 +10,7 @@ from isaaclab_arena.relations.placement_asset import PlaceableAsset from isaaclab_arena.relations.relations import RelationBase -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose, PosePerEnv if TYPE_CHECKING: @@ -27,7 +27,7 @@ class DummyObject(PlaceableAsset): def __init__( self, name: str, - bounding_box: AxisAlignedBoundingBox, + bounding_box: OrientedBoundingBox, initial_pose: Pose | None = None, relations: list[RelationBase] | None = None, collision_mesh: trimesh.Trimesh | None = None, @@ -59,12 +59,12 @@ def _build_reset_event(self) -> EventTermCfg | None: params={"scene_writes": self.layout_pose_to_scene_writes(self.initial_pose)}, ) - def get_bounding_box(self) -> AxisAlignedBoundingBox: + def get_bounding_box(self) -> OrientedBoundingBox: """Get local bounding box (relative to object origin).""" return self.bounding_box def get_corners_aabb(self, pos: torch.Tensor) -> torch.Tensor: - return self.bounding_box.get_corners_at(pos) + return self.bounding_box.translated(pos).get_corners() def is_initial_pose_set(self) -> bool: return self.initial_pose is not None diff --git a/isaaclab_arena/tests/test_bounding_box.py b/isaaclab_arena/tests/test_bounding_box.py index 7d38959e3e..b9173f72cb 100644 --- a/isaaclab_arena/tests/test_bounding_box.py +++ b/isaaclab_arena/tests/test_bounding_box.py @@ -3,203 +3,329 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Tests for AxisAlignedBoundingBox with always-tensor API.""" +"""Tests for oriented bounding-box geometry.""" import math import torch import pytest -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox - - -def test_bounding_box_single_env_properties(): - """Single env: properties return tensors with leading dim 1.""" - aabb = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 4.0, 0.2)) - assert aabb.num_envs == 1 - assert isinstance(aabb.min_point, torch.Tensor) - assert aabb.min_point.shape == (1, 3) - torch.testing.assert_close(aabb.min_point, torch.tensor([[0.0, 0.0, 0.0]])) - torch.testing.assert_close(aabb.max_point[0, 2], torch.tensor(0.2), atol=1e-6, rtol=0) - torch.testing.assert_close(aabb.size, torch.tensor([[2.0, 4.0, 0.2]]), atol=1e-6, rtol=0) - torch.testing.assert_close(aabb.center, torch.tensor([[1.0, 2.0, 0.1]]), atol=1e-6, rtol=0) - assert isinstance(aabb.top_surface_z, torch.Tensor) - assert aabb.top_surface_z.shape == (1,) - torch.testing.assert_close(aabb.top_surface_z, torch.tensor([0.2]), atol=1e-6, rtol=0) - - -def test_bounding_box_single_env_transforms(): - """Single env: translated, scaled, centered, rotated return AABBs with correct values.""" - aabb = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 1.0, 0.5)) - moved = aabb.translated((1.0, 2.0, 0.5)) - torch.testing.assert_close(moved.min_point, torch.tensor([[1.0, 2.0, 0.5]])) - torch.testing.assert_close(moved.max_point, torch.tensor([[3.0, 3.0, 1.0]])) - - scaled = aabb.scaled((2.0, 0.5, 3.0)) - torch.testing.assert_close(scaled.max_point, torch.tensor([[4.0, 0.5, 1.5]])) - - centered = AxisAlignedBoundingBox(min_point=(2.0, 4.0, 0.0), max_point=(4.0, 6.0, 2.0)).centered() - torch.testing.assert_close(centered.center, torch.tensor([[0.0, 0.0, 0.0]]), atol=1e-6, rtol=0) - - rotated = aabb.rotated_90_around_z(1) - torch.testing.assert_close(rotated.min_point, torch.tensor([[-1.0, 0.0, 0.0]]), atol=1e-6, rtol=0) - torch.testing.assert_close(rotated.max_point, torch.tensor([[0.0, 2.0, 0.5]]), atol=1e-6, rtol=0) - - -def test_rotated_around_z_single_angle(): - """90° matches rotated_90_around_z; 45° inflates a centered box to its conservative enclosure.""" - off_origin = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 1.0, 0.5)) - rot90 = off_origin.rotated_around_z(math.pi / 2) - torch.testing.assert_close(rot90.min_point, off_origin.rotated_90_around_z(1).min_point, atol=1e-6, rtol=0) - torch.testing.assert_close(rot90.min_point, torch.tensor([[-1.0, 0.0, 0.0]]), atol=1e-6, rtol=0) - torch.testing.assert_close(rot90.max_point, torch.tensor([[0.0, 2.0, 0.5]]), atol=1e-6, rtol=0) - - # Half-extents (0.2, 0.1); enclosing half-extent = a|cos| + b|sin| = 0.3*cos(45°) on each axis. - centered = AxisAlignedBoundingBox(min_point=(-0.2, -0.1, -0.05), max_point=(0.2, 0.1, 0.05)) - rot45 = centered.rotated_around_z(math.pi / 4) - half = (0.2 + 0.1) * math.cos(math.pi / 4) - torch.testing.assert_close(rot45.min_point, torch.tensor([[-half, -half, -0.05]]), atol=1e-6, rtol=0) - torch.testing.assert_close(rot45.max_point, torch.tensor([[half, half, 0.05]]), atol=1e-6, rtol=0) - - -def test_rotated_around_z_off_center_arbitrary_angle(): - """An off-center box at 30° enclosed by hand-computed corner extents (center shifts, Z fixed).""" - box = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 1.0, 0.5)) - rot = box.rotated_around_z(math.pi / 6) - cos, sin = math.cos(math.pi / 6), math.sin(math.pi / 6) - corners = [(0.0, 0.0), (2.0, 0.0), (2.0, 1.0), (0.0, 1.0)] - xs = [x * cos - y * sin for x, y in corners] - ys = [x * sin + y * cos for x, y in corners] - torch.testing.assert_close(rot.min_point, torch.tensor([[min(xs), min(ys), 0.0]]), atol=1e-6, rtol=0) - torch.testing.assert_close(rot.max_point, torch.tensor([[max(xs), max(ys), 0.5]]), atol=1e-6, rtol=0) - - -def test_rotated_by_quat_encloses_pitched_box(): - """A 90° pitch about Y swaps the X and Z extents, unlike the Z-only rotated_around_z.""" - box = AxisAlignedBoundingBox(min_point=(-0.1, -0.2, -0.5), max_point=(0.1, 0.2, 0.5)) - # Pitch quaternion about Y by 90°, as (x, y, z, w). - pitch_quat = (0.0, math.sin(math.pi / 4), 0.0, math.cos(math.pi / 4)) - rotated = box.rotated_by_quat(pitch_quat) - torch.testing.assert_close(rotated.min_point, torch.tensor([[-0.5, -0.2, -0.1]]), atol=1e-6, rtol=0) - torch.testing.assert_close(rotated.max_point, torch.tensor([[0.5, 0.2, 0.1]]), atol=1e-6, rtol=0) - - -def test_rotated_around_z_batched_angles_broadcasts_single_box(): - """An (M,) angle tensor broadcasts an N=1 box to M enclosing boxes (one per angle).""" - aabb = AxisAlignedBoundingBox(min_point=(-0.2, -0.1, 0.0), max_point=(0.2, 0.1, 0.5)) - angles = torch.tensor([0.0, math.pi / 2]) - rotated = aabb.rotated_around_z(angles) - assert rotated.num_envs == 2 - # Angle 0: unchanged. - torch.testing.assert_close(rotated.min_point[0], torch.tensor([-0.2, -0.1, 0.0]), atol=1e-6, rtol=0) - torch.testing.assert_close(rotated.max_point[0], torch.tensor([0.2, 0.1, 0.5]), atol=1e-6, rtol=0) - # Angle 90°: X/Y extents swap for this origin-centered box. - torch.testing.assert_close(rotated.min_point[1], torch.tensor([-0.1, -0.2, 0.0]), atol=1e-6, rtol=0) - torch.testing.assert_close(rotated.max_point[1], torch.tensor([0.1, 0.2, 0.5]), atol=1e-6, rtol=0) - - -def test_getitem_selects_single_row(): - """Indexing a batched bbox returns the (N=1) box for that row; out-of-range asserts.""" - boxes = AxisAlignedBoundingBox( - min_point=torch.tensor([[0.0, 0.0, 0.0], [1.0, 2.0, 3.0]]), - max_point=torch.tensor([[1.0, 1.0, 1.0], [4.0, 5.0, 6.0]]), - ) - first = boxes[0] - assert first.num_envs == 1 - torch.testing.assert_close(first.min_point, torch.tensor([[0.0, 0.0, 0.0]])) - torch.testing.assert_close(first.max_point, torch.tensor([[1.0, 1.0, 1.0]])) +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox, get_random_pose_within_bounding_box - last = boxes[1] - torch.testing.assert_close(last.min_point, torch.tensor([[1.0, 2.0, 3.0]])) - torch.testing.assert_close(last.max_point, torch.tensor([[4.0, 5.0, 6.0]])) +IDENTITY = (0.0, 0.0, 0.0, 1.0) - with pytest.raises(AssertionError): - _ = boxes[2] +def _yaw(angle: float) -> tuple[float, float, float, float]: + return (0.0, 0.0, math.sin(angle / 2.0), math.cos(angle / 2.0)) -def test_rotated_around_z_mismatched_box_and_angle_counts_raises(): - """Multiple boxes paired with a different count of multiple angles is ambiguous and must assert.""" - boxes = AxisAlignedBoundingBox( - min_point=torch.tensor([[-0.2, -0.1, 0.0], [-0.2, -0.1, 0.0]]), - max_point=torch.tensor([[0.2, 0.1, 0.5], [0.2, 0.1, 0.5]]), - ) - with pytest.raises(AssertionError): - boxes.rotated_around_z(torch.tensor([0.0, math.pi / 2, math.pi])) +def _pitch(angle: float) -> tuple[float, float, float, float]: + return (0.0, math.sin(angle / 2.0), 0.0, math.cos(angle / 2.0)) -def test_bounding_box_single_env_overlaps(): - """Single env: overlaps() returns a (1,) bool tensor; margin widens the check.""" - a = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 1.0)) - b = AxisAlignedBoundingBox(min_point=(0.5, 0.5, 0.0), max_point=(1.5, 1.5, 0.5)) - assert a.overlaps(b).item() is True - c = AxisAlignedBoundingBox(min_point=(1.05, 0.0, 0.0), max_point=(2.0, 1.0, 1.0)) - assert a.overlaps(c).item() is False - assert a.overlaps(c, margin=0.1).item() is True +def test_construction_and_min_max(): + """Construction and min/max conversion preserve box geometry.""" + box = OrientedBoundingBox((1, 2, 3), (0.5, 1, 2), IDENTITY) + assert box.num_envs == 1 + assert box.center.shape == box.half_extents.shape == (1, 3) + assert box.rotation_xyzw.shape == (1, 4) + assert box.center.dtype == box.half_extents.dtype == box.rotation_xyzw.dtype == torch.float32 + from_bounds = OrientedBoundingBox.from_min_max((-1, 0, 2), (3, 4, 8)) + torch.testing.assert_close(from_bounds.center, torch.tensor([[1.0, 2.0, 5.0]])) + torch.testing.assert_close(from_bounds.half_extents, torch.tensor([[2.0, 2.0, 3.0]])) + torch.testing.assert_close(from_bounds.rotation_xyzw, torch.tensor([IDENTITY])) -def test_bounding_box_single_env_get_corners_at(): - """Single env: get_corners_at() returns (1, 8, 3) tensor offset by pos.""" - aabb = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 1.0)) - corners = aabb.get_corners_at() - assert corners.shape == (1, 8, 3) - corners_offset = aabb.get_corners_at(pos=torch.tensor([10.0, 0.0, 0.0])) - torch.testing.assert_close(corners_offset[0, 0], corners[0, 0] + torch.tensor([10.0, 0.0, 0.0])) +def test_constructor_clones_tensor_inputs(): + """Construction clones mutable tensor inputs.""" + center = torch.tensor([1.0, 2.0, 3.0]) + half_extents = torch.tensor([0.5, 1.0, 1.5]) + rotation = torch.tensor(IDENTITY) + box = OrientedBoundingBox(center, half_extents, rotation) + + center.fill_(float("nan")) + half_extents.fill_(-1.0) + rotation.fill_(2.0) + + torch.testing.assert_close(box.center, torch.tensor([[1.0, 2.0, 3.0]])) + torch.testing.assert_close(box.half_extents, torch.tensor([[0.5, 1.0, 1.5]])) + torch.testing.assert_close(box.rotation_xyzw, torch.tensor([IDENTITY])) + + +def test_batch_broadcast_index_invariance_and_translation(): + """Batched boxes broadcast, index, classify, and translate correctly.""" + centers = torch.tensor([[0.0, 0.0, 0.0], [1.0, 2.0, 3.0]]) + box = OrientedBoundingBox(centers, (0.5, 1.0, 1.5), IDENTITY) + assert box.num_envs == 2 + assert box.half_extents.shape == box.rotation_xyzw.shape[:1] + (3,) + assert not box.is_batch_invariant() + torch.testing.assert_close(box[1].center, centers[1:]) + with pytest.raises(AssertionError): + _ = box[2] -def test_bounding_box_multi_env_properties(): - """Multi-env: properties return (N, 3) or (N,) tensors.""" - aabb = AxisAlignedBoundingBox( - min_point=torch.tensor([[0.0, 0.0, 0.0], [1.0, 2.0, 0.0]]), - max_point=torch.tensor([[1.0, 1.0, 1.0], [3.0, 4.0, 0.5]]), + moved = box.translated(torch.tensor([[1.0, 0.0, 0.0], [0.0, -1.0, 2.0]])) + torch.testing.assert_close(moved.center, torch.tensor([[1.0, 0.0, 0.0], [1.0, 1.0, 5.0]])) + invariant = OrientedBoundingBox(torch.zeros(2, 3), torch.ones(2, 3), torch.tensor([IDENTITY])) + assert invariant.is_batch_invariant() + assert invariant.to("cpu").center.device.type == "cpu" + + +def test_off_origin_rotation_and_transform_broadcast(): + """Rigid transforms preserve xyzw composition and row broadcasting.""" + box = OrientedBoundingBox((1.0, 0.0, 0.0), (0.5, 0.25, 0.1), IDENTITY) + rotated = box.rotated_by_quat(_yaw(math.pi / 2)) + torch.testing.assert_close(rotated.center, torch.tensor([[0.0, 1.0, 0.0]]), atol=1e-6, rtol=0) + torch.testing.assert_close(rotated.rotation_xyzw, torch.tensor([_yaw(math.pi / 2)]), atol=1e-6, rtol=0) + + transformed = box.transformed( + torch.tensor([[10.0, 0.0, 1.0], [20.0, 0.0, 2.0]]), + torch.tensor([_yaw(math.pi / 2), _yaw(math.pi)]), ) - assert aabb.num_envs == 2 - assert isinstance(aabb.min_point, torch.Tensor) - assert aabb.size.shape == (2, 3) - torch.testing.assert_close(aabb.size[1], torch.tensor([2.0, 2.0, 0.5])) - torch.testing.assert_close(aabb.center[0], torch.tensor([0.5, 0.5, 0.5])) - assert aabb.top_surface_z.shape == (2,) - torch.testing.assert_close(aabb.top_surface_z, torch.tensor([1.0, 0.5])) - - -def test_bounding_box_multi_env_transforms(): - """Multi-env: translated, scaled, centered, rotated operate per-env.""" - aabb = AxisAlignedBoundingBox( - min_point=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), - max_point=torch.tensor([[2.0, 1.0, 0.5], [3.0, 1.0, 0.5]]), + assert transformed.num_envs == 2 + torch.testing.assert_close( + transformed.center, + torch.tensor([[10.0, 1.0, 1.0], [19.0, 0.0, 2.0]]), + atol=1e-6, + rtol=0, ) - moved = aabb.translated(torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])) - torch.testing.assert_close(moved.min_point[0], torch.tensor([1.0, 0.0, 0.0])) - torch.testing.assert_close(moved.min_point[1], torch.tensor([0.0, 1.0, 0.0])) - rotated = aabb.rotated_90_around_z(1) - torch.testing.assert_close(rotated.min_point[0], torch.tensor([-1.0, 0.0, 0.0])) - torch.testing.assert_close(rotated.max_point[1], torch.tensor([0.0, 3.0, 0.5])) - - centered = aabb.centered() - torch.testing.assert_close(centered.center, torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])) +def test_corners_and_axis_bounds(): + """Corners and projected bounds agree for a yawed box.""" + box = OrientedBoundingBox((1.0, 2.0, 3.0), (2.0, 1.0, 0.5), _yaw(math.pi / 2)) + corners = box.get_corners() + assert corners.shape == (1, 8, 3) + torch.testing.assert_close(corners.mean(dim=1), box.center, atol=1e-6, rtol=0) + expected_min = torch.tensor([[0.0, 0.0, 2.5]]) + expected_max = torch.tensor([[2.0, 4.0, 3.5]]) + minimum, maximum = box.get_axis_aligned_bounds() + torch.testing.assert_close(minimum, expected_min, atol=1e-6, rtol=0) + torch.testing.assert_close(maximum, expected_max, atol=1e-6, rtol=0) + torch.testing.assert_close(corners.amin(dim=1), expected_min, atol=1e-6, rtol=0) + torch.testing.assert_close(corners.amax(dim=1), expected_max, atol=1e-6, rtol=0) + + lower, upper = box.get_bounds_along_axis((2.0, 0.0, 0.0)) + torch.testing.assert_close(lower, torch.tensor([0.0]), atol=1e-6, rtol=0) + torch.testing.assert_close(upper, torch.tensor([2.0]), atol=1e-6, rtol=0) + + +def test_batched_corners_use_each_rotation(): + """Batched corners apply each row's distinct xyzw rotation.""" + box = OrientedBoundingBox( + torch.tensor([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]]), + (2.0, 1.0, 0.5), + torch.tensor([_yaw(math.pi / 2), _pitch(math.pi / 2)]), + ) -def test_bounding_box_multi_env_overlaps(): - """Multi-env vs single-env: overlaps() returns (N,) bool tensor via broadcasting.""" - batched = AxisAlignedBoundingBox( - min_point=torch.tensor([[0.0, 0.0, 0.0], [2.0, 2.0, 0.0]]), - max_point=torch.tensor([[1.0, 1.0, 1.0], [3.0, 3.0, 1.0]]), + corners = box.get_corners() + + torch.testing.assert_close(corners[0].amin(dim=0), torch.tensor([-1.0, -2.0, -0.5]), atol=1e-6, rtol=0) + torch.testing.assert_close(corners[0].amax(dim=0), torch.tensor([1.0, 2.0, 0.5]), atol=1e-6, rtol=0) + torch.testing.assert_close(corners[1].amin(dim=0), torch.tensor([9.5, -1.0, -2.0]), atol=1e-6, rtol=0) + torch.testing.assert_close(corners[1].amax(dim=0), torch.tensor([10.5, 1.0, 2.0]), atol=1e-6, rtol=0) + + +def test_axis_bounds_broadcast_axes(): + """Projection axes broadcast over a single box.""" + box = OrientedBoundingBox((1.0, 2.0, 3.0), (0.5, 1.0, 2.0), IDENTITY) + lower, upper = box.get_bounds_along_axis(torch.eye(3)) + torch.testing.assert_close(lower, torch.tensor([0.5, 1.0, 1.0])) + torch.testing.assert_close(upper, torch.tensor([1.5, 3.0, 5.0])) + + +@pytest.mark.parametrize( + ("rotation", "expected"), + [ + (IDENTITY, True), + (_yaw(math.pi / 2), True), + (_yaw(math.pi), True), + (_yaw(-math.pi / 2), True), + (_pitch(math.pi / 2), True), + (_yaw(math.pi / 4), False), + ], +) +def test_axis_alignment_signed_permutations(rotation, expected): + """Signed axis permutations are classified as axis aligned.""" + box = OrientedBoundingBox((0, 0, 0), (1, 2, 3), rotation) + assert box.is_axis_aligned().item() is expected + + +def test_aabb_overlap_fast_path_touching_and_clearance(): + """The axis-aligned path handles overlap, contact, and clearance.""" + box = OrientedBoundingBox.from_min_max((0, 0, 0), (1, 1, 1)) + overlapping = OrientedBoundingBox.from_min_max((0.75, 0, 0), (1.75, 1, 1)) + touching = OrientedBoundingBox.from_min_max((1, 0, 0), (2, 1, 1)) + separated = OrientedBoundingBox.from_min_max((1.1, 0, 0), (2.1, 1, 1)) + assert box.overlaps(overlapping).item() is True + torch.testing.assert_close(box.penetration(overlapping), torch.tensor([0.25])) + assert box.overlaps(touching).item() is False + torch.testing.assert_close(box.penetration(touching), torch.tensor([0.0])) + assert box.overlaps(separated).item() is False + assert box.overlaps(separated, clearance_m=0.2).item() is True + torch.testing.assert_close(box.penetration(separated, clearance_m=0.2), torch.tensor([0.1]), atol=1e-6, rtol=0) + + +def test_rotated_sat_separation_penetration_and_clearance(): + """The SAT path handles separation, penetration, and clearance.""" + rotation = _yaw(math.pi / 4) + box = OrientedBoundingBox((0, 0, 0), (1.0, 0.2, 0.2), rotation) + perpendicular = torch.tensor([-math.sqrt(0.5), math.sqrt(0.5), 0.0]) + separated = OrientedBoundingBox(perpendicular * 0.5, (1.0, 0.2, 0.2), rotation) + close = OrientedBoundingBox(perpendicular * 0.3, (1.0, 0.2, 0.2), rotation) + touching = OrientedBoundingBox(perpendicular * 0.4, (1.0, 0.2, 0.2), rotation) + + min_a, max_a = box.get_axis_aligned_bounds() + min_b, max_b = separated.get_axis_aligned_bounds() + assert ((max_a > min_b) & (max_b > min_a)).all() + assert box.overlaps(separated).item() is False + assert box.overlaps(close).item() is True + torch.testing.assert_close(box.penetration(close), torch.tensor([0.1]), atol=1e-5, rtol=0) + assert box.overlaps(touching).item() is False + assert box.overlaps(separated, clearance_m=0.2).item() is True + torch.testing.assert_close(box.penetration(separated, clearance_m=0.2), torch.tensor([0.1]), atol=1e-5, rtol=0) + + +def test_identity_fast_path_matches_tiny_angle_sat_path(): + """Identity fast-path depths match tiny-angle SAT depths in mixed batches.""" + centers = torch.tensor([[0.25, 0.1, 0.0], [0.25, 0.1, 0.0]]) + identity_boxes = OrientedBoundingBox(centers, (1.0, 0.8, 0.6), IDENTITY) + mixed_boxes = OrientedBoundingBox(centers, (1.0, 0.8, 0.6), torch.tensor([IDENTITY, _yaw(1e-4)])) + obstacle = OrientedBoundingBox((0.0, 0.0, 0.0), (0.7, 0.5, 0.4), IDENTITY) + + fast = identity_boxes.penetration(obstacle) + mixed = mixed_boxes.penetration(obstacle) + + torch.testing.assert_close(mixed[0], fast[0], atol=1e-6, rtol=0) + torch.testing.assert_close(mixed[1], fast[1], atol=1e-4, rtol=0) + + +def test_overlap_batch_broadcast_and_gradients(): + """Penetration produces finite nonzero gradients in overlapping rows.""" + centers = torch.tensor([[0.0, 0.0, 0.0], [3.0, 0.0, 0.0]], requires_grad=True) + boxes = OrientedBoundingBox(centers, (1, 1, 1), _yaw(math.pi / 4)) + other = OrientedBoundingBox((0.5, 0, 0), (1, 1, 1), IDENTITY) + result = boxes.overlaps(other) + assert result.tolist() == [True, False] + loss = boxes.penetration(other).sum() + loss.backward() + assert centers.grad is not None + assert torch.isfinite(centers.grad).all() + assert torch.count_nonzero(centers.grad[0]).item() > 0 + + +@pytest.mark.parametrize("rotation", [IDENTITY, _yaw(math.pi / 4)]) +def test_concentric_penetration_has_nonzero_center_gradient(rotation): + """Concentric axis-aligned and rotated boxes retain a deterministic escape gradient.""" + center = torch.zeros(3, requires_grad=True) + box = OrientedBoundingBox(center, (1.0, 0.6, 0.4), rotation) + obstacle = OrientedBoundingBox((0.0, 0.0, 0.0), (1.0, 0.6, 0.4), rotation) + + penetration = box.penetration(obstacle) + penetration.sum().backward() + + torch.testing.assert_close(penetration, torch.tensor([0.8]), atol=1e-6, rtol=0) + assert center.grad is not None + assert torch.isfinite(center.grad).all() + assert torch.count_nonzero(center.grad).item() > 0 + + +def test_concentric_cross_axis_tie_breaks_produce_opposite_gradients(): + """Directed ties use one canonical SAT cross axis when box order reverses.""" + from isaaclab.utils.math import quat_from_euler_xyz + + rotation_a = quat_from_euler_xyz( + torch.tensor(0.54799175), + torch.tensor(0.74369037), + torch.tensor(-0.25836289), + ) + rotation_b = quat_from_euler_xyz( + torch.tensor(0.23610878), + torch.tensor(-0.42937827), + torch.tensor(-0.71320212), ) - other = AxisAlignedBoundingBox(min_point=(0.5, 0.5, 0.0), max_point=(1.5, 1.5, 0.5)) - result = batched.overlaps(other) - assert isinstance(result, torch.Tensor) - assert result[0].item() is True - assert result[1].item() is False - - -def test_bounding_box_multi_env_get_corners_at(): - """Multi-env: get_corners_at() returns (N, 8, 3) tensor.""" - aabb = AxisAlignedBoundingBox( - min_point=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), - max_point=torch.tensor([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]), + centers = [torch.zeros(3, requires_grad=True) for _ in range(2)] + boxes = [ + OrientedBoundingBox(centers[0], (2.0, 0.4, 0.15), rotation_a), + OrientedBoundingBox(centers[1], (1.7, 0.3, 0.2), rotation_b), + ] + + forward = boxes[0].penetration( + OrientedBoundingBox(boxes[1].center.detach(), boxes[1].half_extents, boxes[1].rotation_xyzw), + tie_break_sign=1.0, ) - corners = aabb.get_corners_at() - assert corners.shape == (2, 8, 3) - corners_with_pos = aabb.get_corners_at(pos=torch.tensor([[10.0, 0.0, 0.0], [20.0, 0.0, 0.0]])) - torch.testing.assert_close(corners_with_pos[1, 0], corners[1, 0] + torch.tensor([20.0, 0.0, 0.0])) + reverse = boxes[1].penetration( + OrientedBoundingBox(boxes[0].center.detach(), boxes[0].half_extents, boxes[0].rotation_xyzw), + tie_break_sign=-1.0, + ) + (forward + reverse).sum().backward() + + assert centers[0].grad is not None + assert centers[1].grad is not None + assert torch.count_nonzero(centers[0].grad).item() > 0 + torch.testing.assert_close(centers[0].grad, -centers[1].grad, atol=1e-6, rtol=0) + + +def test_penetration_tie_break_does_not_change_nonconcentric_result(): + """Directed tie signs leave non-tied penetration values and gradients unchanged.""" + centers = [torch.tensor([0.3, -0.1, 0.1], requires_grad=True) for _ in range(2)] + obstacle = OrientedBoundingBox((0.0, 0.0, 0.0), (1.0, 0.6, 0.4), _yaw(math.pi / 4)) + values = [] + gradients = [] + + for center, tie_break_sign in zip(centers, (1.0, -1.0)): + value = OrientedBoundingBox(center, (1.0, 0.6, 0.4), _yaw(math.pi / 4)).penetration( + obstacle, tie_break_sign=tie_break_sign + ) + value.sum().backward() + values.append(value) + gradients.append(center.grad) + + torch.testing.assert_close(values[0], values[1]) + torch.testing.assert_close(gradients[0], gradients[1]) + + +@pytest.mark.parametrize( + "constructor", + [ + lambda: OrientedBoundingBox((0, 0), (1, 1, 1), IDENTITY), + lambda: OrientedBoundingBox((0, 0, 0), (-1, 1, 1), IDENTITY), + lambda: OrientedBoundingBox((0, 0, 0), (1, 1, 1), (0, 0, 0, 2)), + lambda: OrientedBoundingBox(torch.zeros(2, 3), torch.ones(3, 3), IDENTITY), + lambda: OrientedBoundingBox.from_min_max((1, 0, 0), (0, 1, 1)), + lambda: OrientedBoundingBox(torch.empty((0, 3)), torch.empty((0, 3)), torch.empty((0, 4))), + lambda: OrientedBoundingBox((float("nan"), 0, 0), (1, 1, 1), IDENTITY), + lambda: OrientedBoundingBox((0, 0, 0), (float("inf"), 1, 1), IDENTITY), + ], +) +def test_invalid_invariants(constructor): + """Invalid dimensions, batches, ranges, and values are rejected.""" + with pytest.raises(AssertionError): + constructor() + + +def test_invalid_clearance_is_rejected(): + """Negative SAT clearance is rejected.""" + box = OrientedBoundingBox((0, 0, 0), (1, 1, 1), IDENTITY) + with pytest.raises(AssertionError, match="Clearance"): + box.penetration(box, clearance_m=-1e-3) + + +def test_zero_projection_axis_is_rejected(): + """A zero projection axis is rejected.""" + box = OrientedBoundingBox((0, 0, 0), (1, 1, 1), IDENTITY) + with pytest.raises(AssertionError, match="non-zero"): + box.get_bounds_along_axis((0.0, 0.0, 0.0)) + + +def test_sampling_reproducibility_and_oriented_membership(): + """Oriented-box sampling is reproducible and remains inside the box.""" + box = OrientedBoundingBox((2.0, 3.0, 4.0), (1.0, 0.5, 0.25), _yaw(math.pi / 2)) + first = get_random_pose_within_bounding_box(box, seed=17) + second = get_random_pose_within_bounding_box(box, seed=17) + assert first == second + position = torch.tensor(first.position_xyz) + delta = position - box.center[0] + local = torch.tensor([delta[1], -delta[0], delta[2]]) + assert (local.abs() <= box.half_extents[0] + 1e-6).all() + assert first.rotation_xyzw == IDENTITY diff --git a/isaaclab_arena/tests/test_embodiment_collision_mesh.py b/isaaclab_arena/tests/test_embodiment_collision_mesh.py index 9b9487d2e6..584a868b2f 100644 --- a/isaaclab_arena/tests/test_embodiment_collision_mesh.py +++ b/isaaclab_arena/tests/test_embodiment_collision_mesh.py @@ -27,7 +27,7 @@ def _test_embodiment_provides_robot_collision_mesh(simulation_app) -> bool: # The mesh and the bounding box describe the same body, so their extents track each other. bbox = emb.get_bounding_box() - bbox_size = (bbox.max_point - bbox.min_point)[0].tolist() + bbox_size = (2.0 * bbox.half_extents[0]).tolist() for mesh_extent, box_extent in zip(extents, bbox_size): assert abs(mesh_extent - box_extent) < 0.2, f"mesh extents {extents} disagree with bbox {bbox_size}" diff --git a/isaaclab_arena/tests/test_face_to.py b/isaaclab_arena/tests/test_face_to.py index f5671a6454..d9f85b6b9a 100644 --- a/isaaclab_arena/tests/test_face_to.py +++ b/isaaclab_arena/tests/test_face_to.py @@ -24,7 +24,7 @@ from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import AtPosition, FaceTo, IsAnchor, RandomAroundSolution, RotateAroundSolution from isaaclab_arena.tests.dummy_object import DummyObject -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange from isaaclab_arena.utils.yaw import wrap_angle_to_pi, yaw_from_quat_xyzw, yaw_toward_positions @@ -36,7 +36,7 @@ def _box(name: str, half_extents: tuple[float, float, float] = (0.1, 0.1, 0.1)) hx, hy, hz = half_extents return DummyObject( name=name, - bounding_box=AxisAlignedBoundingBox(min_point=(-hx, -hy, -hz), max_point=(hx, hy, hz)), + bounding_box=OrientedBoundingBox.from_min_max((-hx, -hy, -hz), (hx, hy, hz)), ) @@ -67,9 +67,8 @@ def _set_solver_results( def _solve( objects: list[ObjectBase], initial_positions: list[dict[ObjectBase, tuple[float, float, float]]], - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None, - env_bboxes_include_yaw: bool = False, - orientations: list[dict[ObjectBase, float]] | None = None, + env_bboxes: dict[ObjectBase, OrientedBoundingBox] | None, + rotations: list[dict[ObjectBase, tuple[float, float, float, float]]] | None = None, collision_objects=None, ) -> list[dict[ObjectBase, tuple[float, float, float]]]: assert len(initial_positions) == len(layouts) @@ -122,12 +121,12 @@ def test_face_to_uses_candidate_positions_for_anchor_and_foreground_targets(): {pair.subject: (0.0, 0.0, 0.0), pair.target: pair.target.get_initial_pose().position_xyz}, {pair.subject: (2.0, -2.0, 0.0), pair.target: (2.0, 0.0, 0.0)}, ] - orientations = [{}, {}] + rotations = [{}, {}] - ObjectPlacer._apply_face_to_orientations(positions, orientations) + ObjectPlacer._apply_face_to_rotations(positions, rotations) - assert orientations[0][pair.subject] == pytest.approx(-math.pi / 2) - assert orientations[1][pair.subject] == pytest.approx(math.pi / 2) + assert yaw_from_quat_xyzw(rotations[0][pair.subject]) == pytest.approx(-math.pi / 2) + assert yaw_from_quat_xyzw(rotations[1][pair.subject]) == pytest.approx(math.pi / 2) def test_coincident_face_to_fails_candidate_validation(): @@ -135,22 +134,22 @@ def test_coincident_face_to_fails_candidate_validation(): subject = _box("subject") subject.add_relation(FaceTo(target)) positions = {subject: (0.0, 0.0, 0.0), target: (0.0, 0.0, 2.0)} - orientations = [{}] + rotations = [{}] - ObjectPlacer._apply_face_to_orientations([positions], orientations) + ObjectPlacer._apply_face_to_rotations([positions], rotations) validation = ObjectPlacer()._validate_candidates( [positions], - [{subject: 0.0}], + rotations, [{obj: obj.get_bounding_box() for obj in positions}], [], )[0] - assert orientations == [{}] + assert rotations == [{}] assert validation.validation_results[PlacementCheck.FACE_TO] is False assert validation.do_all_required_validation_checks_pass() is False -def test_face_to_rebuilds_rotated_footprint_before_validation(): +def test_face_to_rotation_is_applied_directly_during_validation(): target = _box("target") blocker = _box("blocker") subject = _box("subject", half_extents=(1.0, 0.1, 0.1)) @@ -158,17 +157,16 @@ def test_face_to_rebuilds_rotated_footprint_before_validation(): objects = [subject, target, blocker] positions = {subject: (0.0, 0.0, 0.0), target: (0.0, 2.0, 0.0), blocker: (0.0, 0.75, 0.0)} unrotated = {obj: obj.get_bounding_box() for obj in objects} - orientations = [{}] + rotations = [{}] placer = ObjectPlacer() assert placer._validate_candidates([positions], [{}], [unrotated], [])[0].validation_results[ PlacementCheck.NO_OVERLAP ] - ObjectPlacer._apply_face_to_orientations([positions], orientations) - rotated = ObjectPlacer._rotate_candidate_bboxes(objects, unrotated, orientations) - validation = placer._validate_candidates([positions], [orientations[0]], [rotated], [])[0] + ObjectPlacer._apply_face_to_rotations([positions], rotations) + validation = placer._validate_candidates([positions], rotations, [unrotated], [])[0] - assert orientations[0][subject] == pytest.approx(math.pi / 2) + assert yaw_from_quat_xyzw(rotations[0][subject]) == pytest.approx(math.pi / 2) assert validation.validation_results[PlacementCheck.NO_OVERLAP] is False @@ -176,7 +174,7 @@ def test_face_to_suppresses_initial_random_yaw_and_rejects_rotate_marker(): pair = _face_to_pair() placer = ObjectPlacer(ObjectPlacerParams(random_yaw_init=True)) - assert placer._generate_initial_orientations([pair.target, pair.subject], {pair.target}) == {} + assert placer._generate_initial_rotations([pair.target, pair.subject], {pair.target}) == {} pair.subject.add_relation(RotateAroundSolution(yaw_rad=0.5)) with pytest.raises(AssertionError, match="cannot combine FaceTo"): @@ -235,7 +233,7 @@ def test_face_to_allows_reset_rotation_around_facing_yaw(): positions = {pair.target: (0.0, 2.0, 0.0), pair.subject: (0.0, 0.0, 0.0)} anchors, _ = placer._prepare_placement([pair.target, pair.subject]) - placer._apply_poses([positions], anchors, [{pair.subject: math.pi / 2}]) + placer._apply_poses([positions], anchors, [{pair.subject: (0.0, 0.0, 2**-0.5, 2**-0.5)}]) pose = pair.subject.get_initial_pose() assert isinstance(pose, PoseRange) @@ -253,20 +251,54 @@ def test_relation_solver_ignores_face_to_marker(): RelationSolver(RelationSolverParams(max_iters=1, verbose=False)).solve([pair.target, pair.subject], [positions]) -def test_face_to_applies_absolute_world_yaw_with_default_orientation_params(): +def test_coincident_face_to_collision_loss_has_finite_gradients(): + """Coincident FaceTo rows do not poison differentiable OBB collision gradients.""" + from isaaclab_arena.relations.relation_solver_state import RelationSolverState + + pair = _face_to_pair(target_position=(0.0, 0.0, 0.0), subject_half_extents=(0.4, 0.1, 0.1)) + blocker = _box("blocker") + positions = [ + {pair.target: (0.0, 0.0, 0.0), pair.subject: (0.0, 0.0, 0.0), blocker: (0.0, 0.15, 0.0)}, + {pair.target: (0.0, 0.0, 0.0), pair.subject: (-1.0, 0.0, 0.0), blocker: (-1.0, 0.15, 0.0)}, + ] + state = RelationSolverState([pair.target, pair.subject, blocker], positions) + solver = RelationSolver(RelationSolverParams(verbose=False)) + + solver._compute_no_overlap_loss(state).sum().backward() + + assert state.optimizable_positions is not None + gradients = state.optimizable_positions.grad + assert gradients is not None + assert torch.isfinite(gradients).all() + subject_idx = state.optimizable_objects.index(pair.subject) + assert gradients[1, subject_idx].abs().max().item() > 1e-6 + + +def test_face_to_applies_absolute_world_yaw_with_default_orientation_params(monkeypatch): pair = _face_to_pair(target_position=(0.0, 2.0, 0.0)) pair.subject.add_relation(AtPosition(x=0.0, y=0.0, z=0.0)) params = ObjectPlacerParams( solver_params=RelationSolverParams(max_iters=200, verbose=False), max_placement_attempts=1, ) + placer = ObjectPlacer(params) + _set_solver_results( + monkeypatch, + placer, + [{pair.target: (0.0, 2.0, 0.0), pair.subject: (0.0, 0.0, 0.0)}], + ) - (result,) = ObjectPlacer(params).place([pair.target, pair.subject]) + (result,) = placer.place([pair.target, pair.subject]) pose = pair.subject.get_initial_pose() assert isinstance(pose, Pose) - assert result.orientations[pair.subject] == pytest.approx(math.pi / 2) - assert abs(wrap_angle_to_pi(yaw_from_quat_xyzw(pose.rotation_xyzw) - math.pi / 2)) < 1e-5 + expected_yaw = math.atan2( + result.positions[pair.target][1] - result.positions[pair.subject][1], + result.positions[pair.target][0] - result.positions[pair.subject][0], + ) + assert expected_yaw == pytest.approx(math.pi / 2) + assert yaw_from_quat_xyzw(result.rotations[pair.subject]) == pytest.approx(expected_yaw) + assert abs(wrap_angle_to_pi(yaw_from_quat_xyzw(pose.rotation_xyzw) - expected_yaw)) < 1e-5 def test_face_to_applies_independent_multi_env_yaws(monkeypatch): @@ -282,7 +314,7 @@ def test_face_to_applies_independent_multi_env_yaws(monkeypatch): assert isinstance(pose, PosePerEnv) expected = [0.0, math.pi / 2, math.pi, -math.pi / 2] for result, env_pose, expected_yaw in zip(results, pose.poses, expected, strict=True): - assert abs(wrap_angle_to_pi(result.orientations[pair.subject] - expected_yaw)) < 1e-5 + assert abs(wrap_angle_to_pi(yaw_from_quat_xyzw(result.rotations[pair.subject]) - expected_yaw)) < 1e-5 assert abs(wrap_angle_to_pi(yaw_from_quat_xyzw(env_pose.rotation_xyzw) - expected_yaw)) < 1e-5 @@ -329,7 +361,7 @@ def test_face_to_registers_as_binary_graph_relation(): SpatialRelationSpec(kind="face_to", subject="camera") -def test_mesh_validation_receives_final_face_to_yaw(monkeypatch): +def test_mesh_validation_receives_final_face_to_rotation(monkeypatch): pair = _face_to_pair(target_position=(0.0, 1.0, 0.0)) pair.subject.add_relation(AtPosition(x=0.0, y=0.0, z=0.0)) params = ObjectPlacerParams( @@ -338,15 +370,21 @@ def test_mesh_validation_receives_final_face_to_yaw(monkeypatch): apply_positions_to_objects=False, ) placer = ObjectPlacer(params) + _set_solver_results( + monkeypatch, + placer, + [{pair.target: (0.0, 1.0, 0.0), pair.subject: (0.0, 0.0, 0.0)}], + ) received = {} - def _capture_orientations(candidate_positions, env_bboxes, candidate_orientations=None, collision_objects=None): - received.update(candidate_orientations or {}) + def _capture_rotations(candidate_positions, env_bboxes, candidate_rotations=None, collision_objects=None): + received.update(candidate_rotations or {}) return True no_overlap_validator = next(v for v in placer._validators if isinstance(v, NoOverlapValidator)) - monkeypatch.setattr(no_overlap_validator, "_validate_no_overlap_mesh", _capture_orientations) + monkeypatch.setattr(no_overlap_validator, "_validate_no_overlap_mesh", _capture_rotations) (result,) = placer.place([pair.target, pair.subject]) - assert received[pair.subject] == pytest.approx(math.pi / 2) - assert result.orientations[pair.subject] == pytest.approx(math.pi / 2) + expected = (0.0, 0.0, 2**-0.5, 2**-0.5) + assert received[pair.subject] == pytest.approx(expected) + assert result.rotations[pair.subject] == pytest.approx(expected) diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py index 77478a43b1..6182e8ba4f 100644 --- a/isaaclab_arena/tests/test_heterogeneous_placement.py +++ b/isaaclab_arena/tests/test_heterogeneous_placement.py @@ -23,7 +23,7 @@ from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On from isaaclab_arena.tests.dummy_object import DummyObject -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -81,23 +81,24 @@ class HeterogeneousDummyObject(DummyObject): RigidObjectSet's USD machinery. """ - def __init__(self, name: str, bboxes: list[AxisAlignedBoundingBox], **kwargs): + def __init__(self, name: str, bboxes: list[OrientedBoundingBox], **kwargs): super().__init__(name=name, bounding_box=bboxes[0], **kwargs) self._per_env_bboxes = bboxes - def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox: + def get_bounding_box_per_env(self, num_envs: int) -> OrientedBoundingBox: """Return env-specific bbox variants for this test double.""" n_variants = len(self._per_env_bboxes) indices = [i % n_variants for i in range(num_envs)] - min_pts = torch.stack([self._per_env_bboxes[idx].min_point[0] for idx in indices]) - max_pts = torch.stack([self._per_env_bboxes[idx].max_point[0] for idx in indices]) - return AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts) + centers = torch.stack([self._per_env_bboxes[idx].center[0] for idx in indices]) + half_extents = torch.stack([self._per_env_bboxes[idx].half_extents[0] for idx in indices]) + rotations = torch.stack([self._per_env_bboxes[idx].rotation_xyzw[0] for idx in indices]) + return OrientedBoundingBox(center=centers, half_extents=half_extents, rotation_xyzw=rotations) def _make_desk() -> DummyObject: desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) @@ -114,45 +115,45 @@ def test_dummy_object_bbox_per_env_expands_single(): obj = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) per_env = get_bounding_box_per_env(obj, 4) - assert per_env.min_point.shape == (4, 3) - assert per_env.max_point.shape == (4, 3) - assert torch.allclose(per_env.min_point[0], per_env.min_point[3]) + assert per_env.center.shape == (4, 3) + assert per_env.half_extents.shape == (4, 3) + assert torch.allclose(per_env.center[0], per_env.center[3]) def test_per_env_bounding_boxes_formats_solver_and_env_views(): """PerEnvBoundingBoxes should expose solver and one-env bbox formats.""" obj = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.3, 0.4)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.3, 0.4)), ) env_bboxes = build_per_env_bounding_boxes([obj], num_envs=3) solver_bboxes = env_bboxes.get_bounding_boxes_for_solver_candidates(candidates_per_env=2) per_env_bboxes = env_bboxes.get_bounding_boxes_for_all_envs() - assert solver_bboxes[obj].min_point.shape == (6, 3) - assert solver_bboxes[obj].max_point.shape == (6, 3) + assert solver_bboxes[obj].center.shape == (6, 3) + assert solver_bboxes[obj].half_extents.shape == (6, 3) assert len(per_env_bboxes) == 3 - assert per_env_bboxes[1][obj].min_point.shape == (1, 3) - assert torch.allclose(per_env_bboxes[1][obj].max_point[0], torch.tensor([0.2, 0.3, 0.4])) + assert per_env_bboxes[1][obj].center.shape == (1, 3) + assert torch.allclose(per_env_bboxes[1][obj].half_extents[0], torch.tensor([0.1, 0.15, 0.2])) def test_heterogeneous_dummy_returns_different_bboxes(): """HeterogeneousDummyObject should cycle through its member bboxes.""" - small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) - large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3)) + small = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) + large = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3)) obj = HeterogeneousDummyObject(name="set", bboxes=[small, large]) per_env = obj.get_bounding_box_per_env(4) - assert per_env.max_point.shape == (4, 3) + assert per_env.half_extents.shape == (4, 3) # env 0 and 2 should use small; env 1 and 3 should use large - assert torch.allclose(per_env.max_point[0], torch.tensor([0.1, 0.1, 0.1])) - assert torch.allclose(per_env.max_point[1], torch.tensor([0.3, 0.3, 0.3])) + assert torch.allclose(per_env.half_extents[0], torch.tensor([0.05, 0.05, 0.05])) + assert torch.allclose(per_env.half_extents[1], torch.tensor([0.15, 0.15, 0.15])) def test_dummy_object_preserves_constructor_relations(): @@ -162,7 +163,7 @@ def test_dummy_object_preserves_constructor_relations(): anchor_relation = IsAnchor() obj = DummyObject( name="anchor", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)), relations=[anchor_relation], ) @@ -200,7 +201,7 @@ def test_relation_solver_uses_env_bboxes(): desk = _make_desk() box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=0.01)) @@ -212,7 +213,7 @@ def test_relation_solver_uses_env_bboxes(): # Create per-env bboxes with varying sizes across the batch. min_pts = torch.zeros(batch_size, 3) max_pts = torch.stack([torch.tensor([0.1 + 0.05 * i, 0.1 + 0.05 * i, 0.2]) for i in range(batch_size)]) - env_bbox = AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts) + env_bbox = OrientedBoundingBox.from_min_max(min_point=min_pts, max_point=max_pts) solver_params = RelationSolverParams(max_iters=100, convergence_threshold=1e-3, verbose=False) solver = RelationSolver(params=solver_params) @@ -234,8 +235,8 @@ def test_object_placer_heterogeneous_produces_per_env_results(): desk = _make_desk() - small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) - large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3)) + small = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) + large = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3)) hetero_box = HeterogeneousDummyObject(name="hetero_box", bboxes=[small, large]) hetero_box.add_relation(On(desk, clearance_m=0.01)) @@ -263,9 +264,9 @@ def test_object_placer_heterogeneous_z_height_matches_variant(): desk = _make_desk() # "tall" variant: height 0.4 -> bottom at z ~0.11 (desk top 0.1 + clearance 0.01) - tall = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.4)) + tall = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.4)) # "short" variant: height 0.1 -> bottom at z ~0.11 (same clearance) - short = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.1)) + short = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.1)) hetero = HeterogeneousDummyObject(name="hetero", bboxes=[tall, short]) hetero.add_relation(On(desk, clearance_m=0.01)) @@ -299,15 +300,15 @@ def test_mixed_heterogeneous_and_homogeneous_placement(): desk = _make_desk() # A: heterogeneous — small variant in even envs, large in odd envs. - small_a = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) - large_a = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.25, 0.25, 0.25)) + small_a = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) + large_a = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.25, 0.25, 0.25)) obj_a = HeterogeneousDummyObject(name="A", bboxes=[small_a, large_a]) obj_a.add_relation(On(desk, clearance_m=0.01)) # X: homogeneous — same bbox in all envs. obj_x = DummyObject( name="X", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.15)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.15)), ) obj_x.add_relation(On(desk, clearance_m=0.01)) @@ -405,7 +406,7 @@ def test_object_placer_homogeneous_objects_return_multi_env_result(): desk = _make_desk() box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=0.01)) @@ -431,8 +432,8 @@ def test_object_placer_homogeneous_objects_return_multi_env_result(): def _make_hetero_pool_objects(): """Create desk + heterogeneous box for pool tests.""" desk = _make_desk() - small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) - large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3)) + small = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) + large = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3)) hetero = HeterogeneousDummyObject(name="hetero", bboxes=[small, large]) hetero.add_relation(On(desk, clearance_m=0.01)) @@ -737,7 +738,7 @@ def test_pooled_placer_homogeneous_reports_complete_env_rounds(): desk = _make_desk() box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=0.01)) @@ -757,14 +758,14 @@ def test_pooled_placer_homogeneous_reports_complete_env_rounds(): def test_pooled_placer_mixed_heterogeneous_and_homogeneous_objects(): """A pool with mixed object types should match only per-env geometry by env.""" desk = _make_desk() - small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) - large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.25, 0.25, 0.25)) + small = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) + large = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.25, 0.25, 0.25)) hetero = HeterogeneousDummyObject(name="hetero", bboxes=[small, large]) hetero.add_relation(On(desk, clearance_m=0.01)) box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.15)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.15)), ) box.add_relation(On(desk, clearance_m=0.01)) @@ -793,14 +794,14 @@ def test_pooled_placer_multi_set_different_variant_counts(): """ desk = _make_desk() - bottle_small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.2)) - bottle_medium = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.25)) - bottle_large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.3)) + bottle_small = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.2)) + bottle_medium = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.25)) + bottle_large = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.3)) bottles = HeterogeneousDummyObject(name="bottles", bboxes=[bottle_small, bottle_medium, bottle_large]) bottles.add_relation(On(desk, clearance_m=0.01)) - box_small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.1, 0.1)) - box_large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.15, 0.15)) + box_small = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.1, 0.1)) + box_large = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.15, 0.15)) boxes = HeterogeneousDummyObject(name="boxes", bboxes=[box_small, box_large]) boxes.add_relation(On(desk, clearance_m=0.01)) @@ -821,14 +822,14 @@ def test_pooled_placer_multi_set_sample_with_replacement(): """sample_with_replacement with multi-set heterogeneous objects.""" desk = _make_desk() - a_s = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.15)) - a_m = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.2)) - a_l = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.25)) + a_s = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.15)) + a_m = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.2)) + a_l = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.25)) obj_a = HeterogeneousDummyObject(name="A", bboxes=[a_s, a_m, a_l]) obj_a.add_relation(On(desk, clearance_m=0.01)) - b_s = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) - b_l = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.15, 0.12)) + b_s = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) + b_l = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.15, 0.12)) obj_b = HeterogeneousDummyObject(name="B", bboxes=[b_s, b_l]) obj_b.add_relation(On(desk, clearance_m=0.01)) @@ -847,14 +848,14 @@ def test_pooled_placer_multi_set_sample_without_replacement_triggers_refill(): """Exhausting a per-env pool should trigger refill with multi-set objects.""" desk = _make_desk() - v1 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.15)) - v2 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.2)) - v3 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.18)) + v1 = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.15)) + v2 = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.2)) + v3 = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.18)) obj_a = HeterogeneousDummyObject(name="A", bboxes=[v1, v2, v3]) obj_a.add_relation(On(desk, clearance_m=0.01)) - w1 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.1, 0.1)) - w2 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.18, 0.12, 0.12)) + w1 = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.1, 0.1)) + w2 = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.18, 0.12, 0.12)) obj_b = HeterogeneousDummyObject(name="B", bboxes=[w1, w2]) obj_b.add_relation(On(desk, clearance_m=0.01)) @@ -914,8 +915,8 @@ def test_real_rigid_object_set_through_pooled_placer(): can_a = Object(name="can_a", object_type=ObjectType.RIGID, usd_path="/tmp/can_a.usd") can_b = Object(name="can_b", object_type=ObjectType.RIGID, usd_path="/tmp/can_b.usd") - can_a.bounding_box = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.15)) - can_b.bounding_box = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.2)) + can_a.bounding_box = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.15)) + can_b.bounding_box = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.2)) with ( patch("isaaclab_arena.assets.object_set.detect_object_type", return_value=ObjectType.RIGID), diff --git a/isaaclab_arena/tests/test_isaac_sim_debug_draw.py b/isaaclab_arena/tests/test_isaac_sim_debug_draw.py index 9ef56f43c5..101941eefe 100644 --- a/isaaclab_arena/tests/test_isaac_sim_debug_draw.py +++ b/isaaclab_arena/tests/test_isaac_sim_debug_draw.py @@ -5,7 +5,12 @@ """Smoke tests for IsaacSimDebugDraw.""" +import math +from types import SimpleNamespace +from unittest.mock import MagicMock + from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox def smoke_test_debug_draw(simulation_app) -> bool: @@ -31,3 +36,25 @@ def test_isaac_sim_debug_draw_smoke(): """Smoke test: IsaacSimDebugDraw initializes and runs without errors.""" result = run_simulation_app_function(smoke_test_debug_draw) assert result, "IsaacSimDebugDraw smoke test failed" + + +def test_draw_object_bboxes_uses_oriented_corners(): + """Object drawing emits the exact OBB corners in wireframe edge order.""" + from isaaclab_arena.utils.isaac_sim_debug_draw import DEFAULT_COLOR, IsaacSimDebugDraw + + half_yaw = math.pi / 8 + bbox = OrientedBoundingBox((1.0, 2.0, 3.0), (2.0, 1.0, 0.5), (0.0, 0.0, math.sin(half_yaw), math.cos(half_yaw))) + obj = SimpleNamespace(get_world_bounding_box=lambda: bbox) + debug_draw = IsaacSimDebugDraw.__new__(IsaacSimDebugDraw) + debug_draw._draw = MagicMock() + + debug_draw.draw_object_bboxes([obj], thickness=2.0) + + corners = bbox.get_corners()[0].tolist() + edges = [(0, 1), (1, 2), (2, 3), (3, 0), (4, 5), (5, 6), (6, 7), (7, 4), (0, 4), (1, 5), (2, 6), (3, 7)] + debug_draw._draw.draw_lines.assert_called_once_with( + [corners[start] for start, _ in edges], + [corners[end] for _, end in edges], + [DEFAULT_COLOR] * 12, + [2.0] * 12, + ) diff --git a/isaaclab_arena/tests/test_mesh_collision.py b/isaaclab_arena/tests/test_mesh_collision.py index 11b1e4cdde..3f798038be 100644 --- a/isaaclab_arena/tests/test_mesh_collision.py +++ b/isaaclab_arena/tests/test_mesh_collision.py @@ -20,7 +20,7 @@ from isaaclab_arena.relations.relations import IsAnchor, On from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache, greedy_sphere_decomposition from isaaclab_arena.tests.dummy_object import DummyObject -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose try: @@ -34,6 +34,14 @@ requires_warp = pytest.mark.skipif(not _WARP_AVAILABLE, reason="Warp not available") +def _quat(roll: float = 0.0, pitch: float = 0.0, yaw: float = 0.0) -> tuple[float, float, float, float]: + """Return an xyzw quaternion for Euler XYZ angles.""" + from isaaclab.utils.math import quat_from_euler_xyz + + quat = quat_from_euler_xyz(torch.tensor(roll), torch.tensor(pitch), torch.tensor(yaw)) + return tuple(quat.tolist()) + + # Unit tests @@ -41,7 +49,7 @@ def _make_cylinder(name: str, radius: float = 0.033, height: float = 0.1) -> Dum mesh = trimesh.creation.cylinder(radius=radius, height=height, sections=32) return DummyObject( name=name, - bounding_box=AxisAlignedBoundingBox( + bounding_box=OrientedBoundingBox.from_min_max( min_point=(-radius, -radius, -height / 2), max_point=(radius, radius, height / 2), ), @@ -53,7 +61,7 @@ def _make_box_obj(name: str, sx: float, sy: float, sz: float) -> DummyObject: mesh = trimesh.creation.box(extents=(sx, sy, sz)) return DummyObject( name=name, - bounding_box=AxisAlignedBoundingBox( + bounding_box=OrientedBoundingBox.from_min_max( min_point=(-sx / 2, -sy / 2, -sz / 2), max_point=(sx / 2, sy / 2, sz / 2), ), @@ -65,7 +73,7 @@ def _make_table() -> DummyObject: mesh = trimesh.creation.box(extents=(1.0, 1.0, 0.05)) table = DummyObject( name="table", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.5, -0.5, -0.025), max_point=(0.5, 0.5, 0.025)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.5, -0.5, -0.025), max_point=(0.5, 0.5, 0.025)), collision_mesh=mesh, ) table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) @@ -75,7 +83,7 @@ def _make_table() -> DummyObject: def _env_bboxes_for( positions: dict[DummyObject, tuple[float, float, float]], -) -> dict[DummyObject, AxisAlignedBoundingBox]: +) -> dict[DummyObject, OrientedBoundingBox]: """Return default bboxes keyed by the positioned objects.""" return {obj: obj.get_bounding_box() for obj in positions} @@ -102,33 +110,42 @@ def test_object_placer_aabb_proxy_uses_candidate_bbox(): """Mesh validation builds AABB proxies from the candidate bbox.""" from isaaclab_arena.relations.placement_validators import NoOverlapValidator - bbox = AxisAlignedBoundingBox(min_point=(-0.2, -0.1, -0.05), max_point=(0.2, 0.1, 0.05)) - proxy = NoOverlapValidator._collision_mesh_or_aabb_proxy(None, bbox) + bbox = OrientedBoundingBox.from_min_max(min_point=(-0.2, -0.1, -0.05), max_point=(0.2, 0.1, 0.05)) + proxy = NoOverlapValidator._collision_mesh_or_bbox_proxy(None, bbox) np.testing.assert_allclose(proxy.extents, [0.4, 0.2, 0.1], atol=1e-6) -def test_effective_yaw_ignores_placed_initial_pose_unless_allowed(): - """Placed non-anchors do not inherit initial_pose yaw unless the caller explicitly allows pose yaw.""" - from isaaclab_arena.relations.placement_validators import NoOverlapValidator - - obj = _make_box_obj("placed", sx=0.1, sy=0.02, sz=0.05) - obj.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.7071068, 0.7071068))) - - assert NoOverlapValidator._effective_yaw(obj, orientations=None, use_pose_yaw=False) == 0.0 - assert NoOverlapValidator._effective_yaw(obj, orientations=None, use_pose_yaw=True) > 1.5 - - -def test_mesh_broadphase_rotates_bbox_about_object_origin(): - """Mesh broadphase bbox rotation matches AxisAlignedBoundingBox origin-frame semantics.""" - from isaaclab_arena.relations.no_overlap_mesh import _rotate_bbox_extents - - bbox = AxisAlignedBoundingBox(min_point=(0.1, -0.01, -0.02), max_point=(0.3, 0.01, 0.02)) - expected = bbox.rotated_around_z(torch.tensor([math.pi / 2])) - min_point, max_point = _rotate_bbox_extents(bbox.min_point, bbox.max_point, torch.tensor([math.pi / 2])) +@pytest.mark.parametrize( + ("source_rpy", "target_rpy"), + [ + ((0.4, -0.3, 0.0), (0.0, 0.0, 0.0)), + ((0.0, 0.0, 0.0), (-0.2, 0.5, 0.0)), + ((0.4, -0.3, 0.2), (-0.2, 0.5, -0.1)), + ], +) +def test_full_quaternion_relative_transform(source_rpy, target_rpy): + """Source-only, target-only, and combined roll/pitch use the full relative transform.""" + from isaaclab.utils.math import quat_from_euler_xyz + + from isaaclab_arena.relations.no_overlap_mesh import transform_points_between_frames + + source_rotation = quat_from_euler_xyz(*(torch.tensor(value) for value in source_rpy)) + target_rotation = quat_from_euler_xyz(*(torch.tensor(value) for value in target_rpy)) + centers = torch.tensor([[0.1, -0.02, 0.03]]) + source_position = torch.tensor([0.3, -0.1, 0.2]) + target_position = torch.tensor([-0.2, 0.4, 0.1]) + + transformed = transform_points_between_frames( + centers, source_position, source_rotation, target_position, target_rotation + ) + from isaaclab.utils.math import quat_apply, quat_apply_inverse - torch.testing.assert_close(min_point, expected.min_point) - torch.testing.assert_close(max_point, expected.max_point) + expected = quat_apply_inverse( + target_rotation.unsqueeze(0), + quat_apply(source_rotation.unsqueeze(0), centers) + source_position - target_position, + ) + torch.testing.assert_close(transformed, expected) @requires_warp @@ -158,13 +175,13 @@ def test_warp_mesh_raw_mesh_flag_skips_convex_hull(monkeypatch): ) raw_obj = DummyObject( "raw_background", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.0)), collision_mesh=mesh, ) raw_obj.repair_collision_mesh_non_watertight = False normal_obj = DummyObject( "normal_background", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.0)), collision_mesh=mesh, ) @@ -180,12 +197,9 @@ def fail_convex_hull(self): def _batched_aabb_loss(strategy, clearance_m, child_pos, child_bbox, parent_world_bbox): - """Helper: compute single-pair AABB loss via compute_loss_batched.""" - subject_min = (child_pos + child_bbox.min_point).unsqueeze(0).unsqueeze(0) - subject_max = (child_pos + child_bbox.max_point).unsqueeze(0).unsqueeze(0) - obstacle_min = parent_world_bbox.min_point.unsqueeze(0) - obstacle_max = parent_world_bbox.max_point.unsqueeze(0) - loss = strategy.compute_loss_batched(clearance_m, subject_min, subject_max, obstacle_min, obstacle_max) + """Helper: compute one directed OBB loss via compute_loss_batched.""" + penetration = child_bbox.translated(child_pos).penetration(parent_world_bbox, clearance_m).unsqueeze(0) + loss = strategy.compute_loss_batched(penetration) return loss.squeeze() @@ -343,40 +357,21 @@ def test_anchor_with_rotate_around_solution_rejected(): placer.place([table, child]) -@requires_warp -def test_centers_in_target_frame_applies_both_yaws(): - """Net yaw = source - target; equal yaws cancel out.""" +def test_relative_transform_identity_and_equal_yaw_regressions(): + """Identity passes through and equal source/target yaw cancels with zero offset.""" + from isaaclab.utils.math import quat_from_euler_xyz - from isaaclab_arena.relations.placement_validators import NoOverlapValidator + from isaaclab_arena.relations.no_overlap_mesh import transform_points_between_frames - src = DummyObject( - "src", - 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=(0.2, 0.2, 0.2)), - ) - tgt = DummyObject( - "tgt", - 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=(0.2, 0.2, 0.2)), - ) centers = torch.tensor([[0.10, 0.0, 0.0]]) - src_pos = torch.tensor([0.0, 0.0, 0.0]) - tgt_pos = torch.tensor([0.0, 0.0, 0.0]) + zero = torch.zeros(3) + identity = torch.tensor([0.0, 0.0, 0.0, 1.0]) + torch.testing.assert_close(transform_points_between_frames(centers, zero, identity, zero, identity), centers) - # No orientations: pass-through - result = NoOverlapValidator._centers_in_target_frame(centers, src, tgt, src_pos, tgt_pos, None) - assert torch.allclose(result, centers, atol=1e-6) - - # Source yaw=pi/2, target yaw=0: net rotation = pi/2 - result = NoOverlapValidator._centers_in_target_frame(centers, src, tgt, src_pos, tgt_pos, {src: math.pi / 2}) - assert abs(result[0, 0].item()) < 1e-5 - assert abs(result[0, 1].item() - 0.10) < 1e-5 - - # Both at same yaw: net rotation = 0, centers unchanged (offset is zero here) - result = NoOverlapValidator._centers_in_target_frame( - centers, src, tgt, src_pos, tgt_pos, {src: math.pi / 2, tgt: math.pi / 2} + yaw = quat_from_euler_xyz(torch.tensor(0.0), torch.tensor(0.0), torch.tensor(math.pi / 2)) + torch.testing.assert_close( + transform_points_between_frames(centers, zero, yaw, zero, yaw), centers, atol=1e-6, rtol=1e-6 ) - assert torch.allclose(result, centers, atol=1e-5) @requires_warp @@ -435,7 +430,7 @@ def test_validate_placement_mesh_mode_rejects_aabb_foreground_background_overlap table = _make_table() box = DummyObject( "aabb_only_box", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.05, -0.05, -0.05), max_point=(0.05, 0.05, 0.05)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.05, -0.05, -0.05), max_point=(0.05, 0.05, 0.05)), ) box.add_relation(On(table)) background = _make_box_obj("mesh_background", 0.2, 0.2, 0.2) @@ -459,7 +454,7 @@ def test_validate_placement_mesh_mode_rejects_aabb_foreground_background_overlap @requires_warp def test_validate_no_overlap_mesh_sentinel_fails(monkeypatch): - """A sentinel SDF (no resolvable face) must fail validation, not certify collision-free.""" + """CPU validation rejects unsupported sentinel geometry with the solver message.""" from isaaclab_arena.relations import warp_sdf_kernels from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.placement_validators import NoOverlapValidator @@ -479,7 +474,7 @@ def test_validate_no_overlap_mesh_sentinel_fails(monkeypatch): env_bboxes = _env_bboxes_for(positions) assert validator._validate_no_overlap_mesh(positions, env_bboxes) - # Force every query to hit the sentinel; the same separated layout must now fail. + # Force every query to hit the sentinel; the same separated layout is unsupported. from isaaclab_arena.relations import placement_validators as _pv_mod real_mesh_sdf = warp_sdf_kernels.mesh_sdf @@ -489,7 +484,11 @@ def fake_sdf(points, mesh): monkeypatch.setattr(warp_sdf_kernels, "mesh_sdf", fake_sdf) monkeypatch.setattr(_pv_mod, "mesh_sdf", fake_sdf) - assert not validator._validate_no_overlap_mesh(positions, env_bboxes) + with pytest.raises( + AssertionError, + match="MESH collision query could not resolve a target face; the promised collision geometry is unsupported", + ): + validator._validate_no_overlap_mesh(positions, env_bboxes) @requires_warp @@ -503,7 +502,7 @@ def test_validate_no_overlap_mesh_respects_anchor_yaw(): anchor_mesh = trimesh.creation.box(extents=(0.2, 0.02, 0.05)) anchor = DummyObject( "anchor", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.1, -0.01, -0.025), max_point=(0.1, 0.01, 0.025)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.1, -0.01, -0.025), max_point=(0.1, 0.01, 0.025)), collision_mesh=anchor_mesh, ) sz = math.sin(math.pi / 4) @@ -623,6 +622,47 @@ def test_broadphase_does_not_skip_overlapping_pairs(): assert loss > 0.0, "Overlapping objects should produce nonzero loss" +@requires_warp +def test_broadphase_rotates_off_origin_base_obb(): + """A rotated off-origin base OBB must keep a genuinely overlapping pair active.""" + table = _make_table() + table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, -1.0))) + source_mesh = trimesh.creation.box(extents=(0.08, 0.02, 0.02)) + source_mesh.apply_translation((0.2, 0.0, 0.0)) + source = DummyObject( + "off_origin_source", + bounding_box=OrientedBoundingBox( + center=(0.2, 0.0, 0.0), + half_extents=(0.04, 0.01, 0.01), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), + ), + collision_mesh=source_mesh, + ) + target = _make_box_obj("target", sx=0.05, sy=0.05, sz=0.05) + initial = [{table: (0.0, 0.0, -1.0), source: (0.0, 0.0, 0.0), target: (0.0, 0.2, 0.0)}] + + solver = RelationSolver(params=RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=0, verbose=False)) + solver.solve([table, source, target], initial, rotations=[{source: _quat(yaw=math.pi / 2)}]) + + assert solver.last_loss_per_env[0].item() > 0.0 + + +@requires_warp +def test_fixed_mesh_obstacle_roll_pitch_is_supported(): + """Fixed mesh obstacles with roll/pitch are transformed, not rejected or yaw-projected.""" + table = _make_table() + table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, -1.0))) + source = _make_cylinder("source", radius=0.012, height=0.02) + obstacle = _make_box_obj("pitched_obstacle", sx=0.2, sy=0.02, sz=0.05) + obstacle.set_initial_pose(Pose(rotation_xyzw=_quat(roll=0.2, pitch=math.pi / 2))) + initial = [{table: (0.0, 0.0, -1.0), source: (0.0, 0.0, 0.06)}] + + solver = RelationSolver(params=RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=0, verbose=False)) + solver.solve([table, source], initial, collision_objects=[obstacle]) + + assert solver.last_loss_per_env[0].item() > 0.0 + + @requires_warp def test_object_collision_mode_can_force_bbox_in_mesh_solver(): """Objects can opt out of mesh collision even when the solver default is MESH.""" @@ -667,7 +707,7 @@ def test_mixed_mesh_aabb_uses_per_env_bbox_proxy(): table = _make_table() source = DummyObject( "source", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.01, -0.01, -0.01), max_point=(0.01, 0.01, 0.01)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.01, -0.01, -0.01), max_point=(0.01, 0.01, 0.01)), ) target = _make_box_obj("target", sx=0.1, sy=0.1, sz=0.1) target.collision_mode = CollisionMode.MESH @@ -675,7 +715,7 @@ def test_mixed_mesh_aabb_uses_per_env_bbox_proxy(): initial = [{table: (0.0, 0.0, 0.0), source: (0.0, 0.0, 0.2), target: (0.08, 0.0, 0.2)}] env_bboxes = { table: table.get_bounding_box(), - source: AxisAlignedBoundingBox(min_point=(-0.06, -0.06, -0.01), max_point=(0.06, 0.06, 0.01)), + source: OrientedBoundingBox.from_min_max(min_point=(-0.06, -0.06, -0.01), max_point=(0.06, 0.06, 0.01)), target: target.get_bounding_box(), } @@ -683,6 +723,16 @@ def test_mixed_mesh_aabb_uses_per_env_bbox_proxy(): solver.solve([table, source, target], initial, env_bboxes=env_bboxes) assert solver._mesh_cache is not None + directions = { + (subject.name, obstacle.name) + for subject, obstacle in zip( + solver._mesh_cache.pair_subject_objs, + solver._mesh_cache.pair_obstacle_objs, + strict=True, + ) + } + assert ("source", "target") in directions + assert ("target", "source") in directions assert solver.last_loss_per_env[0].item() > 0.0 @@ -690,9 +740,10 @@ def test_mixed_mesh_aabb_uses_per_env_bbox_proxy(): def test_mixed_mesh_aabb_varying_proxy_uses_aabb_fallback(): """Varying per-env AABB proxies stay on the AABB collision path.""" table = _make_table() + table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, -1.0))) source = DummyObject( "source", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.01, -0.01, -0.01), max_point=(0.01, 0.01, 0.01)), + bounding_box=OrientedBoundingBox.from_min_max(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 @@ -703,7 +754,7 @@ def test_mixed_mesh_aabb_varying_proxy_uses_aabb_fallback(): ] env_bboxes = { table: table.get_bounding_box(), - source: AxisAlignedBoundingBox( + source: OrientedBoundingBox.from_min_max( min_point=torch.tensor([[-0.01, -0.01, -0.01], [-0.3, -0.3, -0.01]]), max_point=torch.tensor([[0.01, 0.01, 0.01], [0.3, 0.3, 0.01]]), ), @@ -722,37 +773,66 @@ def test_mixed_mesh_aabb_varying_proxy_uses_aabb_fallback(): @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.""" + """A base OBB proxy receives its candidate root rotation exactly once.""" 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.2, -0.01, -0.025), max_point=(0.2, 0.01, 0.025)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.2, -0.01, -0.025), max_point=(0.2, 0.01, 0.025)), ) source.collision_mode = CollisionMode.BBOX 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.0, 0.15, 0.0)} - env_bboxes = { - source: source.get_bounding_box().rotated_around_z(torch.tensor([math.pi / 2])), - target: target.get_bounding_box(), - } + env_bboxes = {source: source.get_bounding_box(), target: target.get_bounding_box()} + validator = NoOverlapValidator( + ObjectPlacerParams(solver_params=RelationSolverParams(collision_mode=CollisionMode.BBOX, verbose=False)) + ) + + assert not validator._validate_no_overlap_mesh(positions, env_bboxes, rotations={source: _quat(yaw=math.pi / 2)}) + + +@requires_warp +def test_mixed_mesh_bbox_validation_checks_both_directions(monkeypatch): + """Validation mirrors the solver's two directed mixed mesh/proxy checks.""" + from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams + from isaaclab_arena.relations.placement_validators import NoOverlapValidator + + source = DummyObject( + "source", + bounding_box=OrientedBoundingBox.from_min_max((-0.05, -0.05, -0.05), (0.05, 0.05, 0.05)), + ) + source.collision_mode = CollisionMode.BBOX + target = _make_box_obj("target", sx=0.1, sy=0.1, sz=0.1) + target.collision_mode = CollisionMode.MESH + positions = {source: (0.0, 0.0, 0.0), target: (0.0, 0.0, 0.0)} + env_bboxes = {obj: obj.get_bounding_box() for obj in positions} validator = NoOverlapValidator( ObjectPlacerParams(solver_params=RelationSolverParams(collision_mode=CollisionMode.BBOX, verbose=False)) ) + directions: list[tuple[str, str]] = [] - assert not validator._validate_no_overlap_mesh(positions, env_bboxes, orientations={source: math.pi / 2}) + def _capture_direction(*args): + direction = (args[0].name, args[5].name) + directions.append(direction) + return direction == ("target", "source") + + monkeypatch.setattr(validator, "_spheres_penetrate_mesh", _capture_direction) + + assert not validator._validate_no_overlap_mesh(positions, env_bboxes) + assert directions == [("source", "target"), ("target", "source")] @requires_warp def test_yawed_aabb_proxy_solver_loss_is_not_double_rotated(): - """Solver mixed mesh/AABB loss uses yaw-expanded proxy bboxes without rotating them again.""" + """Solver mixed mesh/OBB loss transforms the asset-local proxy exactly once.""" table = _make_table() + table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, -1.0))) source = DummyObject( "source", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.2, -0.01, -0.025), max_point=(0.2, 0.01, 0.025)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.2, -0.01, -0.025), max_point=(0.2, 0.01, 0.025)), ) source.collision_mode = CollisionMode.BBOX target = _make_box_obj("target", sx=0.05, sy=0.05, sz=0.05) @@ -760,7 +840,7 @@ def test_yawed_aabb_proxy_solver_loss_is_not_double_rotated(): initial = [{table: (0.0, 0.0, -1.0), source: (0.0, 0.0, 0.0), target: (0.0, 0.15, 0.0)}] env_bboxes = { table: table.get_bounding_box(), - source: source.get_bounding_box().rotated_around_z(torch.tensor([math.pi / 2])), + source: source.get_bounding_box(), target: target.get_bounding_box(), } @@ -769,29 +849,34 @@ def test_yawed_aabb_proxy_solver_loss_is_not_double_rotated(): [table, source, target], initial, env_bboxes=env_bboxes, - env_bboxes_include_yaw=True, - orientations=[{source: math.pi / 2}], + rotations=[{source: _quat(yaw=math.pi / 2)}], ) assert solver.last_loss_per_env[0].item() > 0.0 @requires_warp -def test_yawed_aabb_proxy_solver_loss_rotates_unexpanded_bbox(): - """Direct solver calls rotate AABB proxy spheres when bboxes are not pre-expanded.""" +def test_rolled_meshless_proxy_solver_loss_uses_full_rotation(): + """A meshless box proxy uses source roll, not a yaw projection.""" table = _make_table() + table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, -1.0))) source = DummyObject( "source", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.2, -0.01, -0.025), max_point=(0.2, 0.01, 0.025)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.2, -0.01, -0.025), max_point=(0.2, 0.01, 0.025)), ) source.collision_mode = CollisionMode.BBOX target = _make_box_obj("target", sx=0.05, sy=0.05, sz=0.05) target.collision_mode = CollisionMode.MESH - initial = [{table: (0.0, 0.0, -1.0), source: (0.0, 0.0, 0.0), target: (0.0, 0.15, 0.0)}] + initial = [{table: (0.0, 0.0, -1.0), source: (0.0, 0.0, 0.0), target: (0.0, 0.0, 0.15)}] env_bboxes = {table: table.get_bounding_box(), source: source.get_bounding_box(), target: target.get_bounding_box()} solver = RelationSolver(params=RelationSolverParams(collision_mode=CollisionMode.BBOX, max_iters=0, verbose=False)) - solver.solve([table, source, target], initial, env_bboxes=env_bboxes, orientations=[{source: math.pi / 2}]) + solver.solve( + [table, source, target], + initial, + env_bboxes=env_bboxes, + rotations=[{source: _quat(pitch=-math.pi / 2)}], + ) assert solver.last_loss_per_env[0].item() > 0.0 @@ -804,7 +889,7 @@ def test_validate_no_overlap_mesh_respects_yawed_collision_object(): source = DummyObject( "source", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.01, -0.01, -0.01), max_point=(0.01, 0.01, 0.01)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.01, -0.01, -0.01), max_point=(0.01, 0.01, 0.01)), ) source.collision_mode = CollisionMode.BBOX obstacle = _make_box_obj("obstacle", sx=0.2, sy=0.02, sz=0.05) @@ -825,9 +910,10 @@ def test_validate_no_overlap_mesh_respects_yawed_collision_object(): def test_solver_mesh_loss_broadphase_respects_yawed_collision_object(): """Fixed mesh obstacle bboxes still rotate when placed-object bboxes are yaw-expanded.""" table = _make_table() + table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, -1.0))) source = DummyObject( "source", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.01, -0.01, -0.01), max_point=(0.01, 0.01, 0.01)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.01, -0.01, -0.01), max_point=(0.01, 0.01, 0.01)), ) source.collision_mode = CollisionMode.BBOX obstacle = _make_box_obj("obstacle", sx=0.2, sy=0.02, sz=0.05) @@ -839,7 +925,7 @@ def test_solver_mesh_loss_broadphase_respects_yawed_collision_object(): env_bboxes = {table: table.get_bounding_box(), source: source.get_bounding_box()} solver = RelationSolver(params=RelationSolverParams(collision_mode=CollisionMode.BBOX, max_iters=0, verbose=False)) - solver.solve([table, source], initial, env_bboxes=env_bboxes, orientations=[{}], collision_objects=[obstacle]) + solver.solve([table, source], initial, env_bboxes=env_bboxes, rotations=[{}], collision_objects=[obstacle]) assert solver.last_loss_per_env[0].item() > 0.0 @@ -912,23 +998,21 @@ def test_solver_target_only_yaw(): # With target rotated 90°, the 0.2/2=0.1 half-extent now spans Y → collision. initial = [{table: (0.0, 0.0, 0.0), target: (0.0, 0.0, 0.05), child: (0.0, 0.03, 0.05)}] - solver_no_rot = RelationSolver( - params=RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=0, verbose=False) - ) - solver_no_rot.solve([table, target, child], initial, orientations=None) + params = RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=0, verbose=False) + solver_no_rot = RelationSolver(params=params) + solver_no_rot.solve([table, target, child], initial) loss_no_rot = solver_no_rot.last_loss_per_env[0].item() # Target rotated 90° around Z: child is now inside target's mesh - orientations_rotated = [{target: math.pi / 2, child: 0.0}] - solver_rot = RelationSolver( - params=RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=0, verbose=False) - ) - solver_rot.solve([table, target, child], initial, orientations=orientations_rotated) + target.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.05), rotation_xyzw=_quat(yaw=math.pi / 2))) + solver_rot = RelationSolver(params=params) + solver_rot.solve([table, target, child], initial) loss_rot = solver_rot.last_loss_per_env[0].item() + minimum_loss_delta = 1e-4 * params.collision_loss_slope assert ( - loss_rot > loss_no_rot + 1.0 - ), f"Target yaw=90° should dramatically increase collision loss (got {loss_rot:.2f} vs {loss_no_rot:.2f})" + loss_rot > loss_no_rot + minimum_loss_delta + ), f"Target yaw=90° should increase collision loss (got {loss_rot:.2f} vs {loss_no_rot:.2f})" @requires_warp @@ -939,7 +1023,7 @@ def test_anchor_initial_pose_yaw_affects_collision(): target_mesh = trimesh.creation.box(extents=(0.2, 0.02, 0.05)) target = DummyObject( "target", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.1, -0.01, -0.025), max_point=(0.1, 0.01, 0.025)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.1, -0.01, -0.025), max_point=(0.1, 0.01, 0.025)), collision_mesh=target_mesh, ) # Bake 90° Z-yaw into initial_pose (not via orientations dict) @@ -954,27 +1038,27 @@ def test_anchor_initial_pose_yaw_affects_collision(): # Child at Y=0.02: outside unrotated target (half-width=0.01), inside rotated (half-length=0.1) initial = [{table: (0.0, 0.0, 0.0), target: (0.0, 0.0, 0.05), child: (0.0, 0.02, 0.05)}] - solver = RelationSolver(params=RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=0, verbose=False)) - solver.solve([table, target, child], initial, orientations=None) + params = RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=0, verbose=False) + solver = RelationSolver(params=params) + solver.solve([table, target, child], initial) loss_yawed = solver.last_loss_per_env[0].item() # Same geometry with identity anchor — should have lower loss target_id = DummyObject( "target_id", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.1, -0.01, -0.025), max_point=(0.1, 0.01, 0.025)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.1, -0.01, -0.025), max_point=(0.1, 0.01, 0.025)), collision_mesh=target_mesh, ) target_id.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.05))) target_id.add_relation(IsAnchor()) initial_id = [{table: (0.0, 0.0, 0.0), target_id: (0.0, 0.0, 0.05), child: (0.0, 0.02, 0.05)}] - solver_id = RelationSolver( - params=RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=0, verbose=False) - ) - solver_id.solve([table, target_id, child], initial_id, orientations=None) + solver_id = RelationSolver(params=params) + solver_id.solve([table, target_id, child], initial_id) loss_identity = solver_id.last_loss_per_env[0].item() - assert loss_yawed > loss_identity + 1.0, ( + minimum_loss_delta = 1e-4 * params.collision_loss_slope + assert loss_yawed > loss_identity + minimum_loss_delta, ( "Yawed anchor (from initial_pose) should produce higher collision " f"(got {loss_yawed:.2f} vs identity {loss_identity:.2f})" ) @@ -1031,12 +1115,12 @@ def test_broadphase_does_not_falsely_cull_yawed_elongated_pair(): b_mesh = trimesh.creation.box(extents=(0.4, 0.02, 0.05)) a = DummyObject( "a", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.2, -0.01, -0.025), max_point=(0.2, 0.01, 0.025)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.2, -0.01, -0.025), max_point=(0.2, 0.01, 0.025)), collision_mesh=a_mesh, ) b = DummyObject( "b", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.2, -0.01, -0.025), max_point=(0.2, 0.01, 0.025)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.2, -0.01, -0.025), max_point=(0.2, 0.01, 0.025)), collision_mesh=b_mesh, ) a.add_relation(On(table)) @@ -1046,10 +1130,10 @@ def test_broadphase_does_not_falsely_cull_yawed_elongated_pair(): # but if we yaw both by pi/2 the unrotated AABB (width 0.02) would be separated. yaw = math.pi / 2 initial = [{table: (0.0, 0.0, 0.0), a: (0.0, 0.0, 0.05), b: (0.05, 0.0, 0.05)}] - orientations = [{a: yaw, b: yaw}] + rotations = [{a: _quat(yaw=yaw), b: _quat(yaw=yaw)}] solver = RelationSolver(params=RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=0, verbose=False)) - solver.solve([table, a, b], initial, orientations=orientations) + solver.solve([table, a, b], initial, rotations=rotations) loss = solver.last_loss_per_env[0].item() assert loss > 0.0, "Broadphase must not falsely cull yawed elongated pairs that genuinely collide" @@ -1060,7 +1144,7 @@ def test_mesh_mode_queries_aabb_subject_against_mesh_background(): table = _make_table() box = DummyObject( "aabb_only_box", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.05, -0.05, -0.05), max_point=(0.05, 0.05, 0.05)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.05, -0.05, -0.05), max_point=(0.05, 0.05, 0.05)), ) box.add_relation(On(table)) background = _make_box_obj("mesh_background", 0.2, 0.2, 0.2) @@ -1089,7 +1173,7 @@ def test_mesh_mode_scores_background_collision_object(): table = _make_table() box = DummyObject( "aabb_only_box", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.05, -0.05, -0.05), max_point=(0.05, 0.05, 0.05)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.05, -0.05, -0.05), max_point=(0.05, 0.05, 0.05)), ) box.add_relation(On(table)) background = FixedCollisionObject(trimesh.creation.box(extents=(0.2, 0.2, 0.2))) @@ -1118,7 +1202,7 @@ def test_mesh_mode_scores_mixed_mesh_aabb_placed_pair(): mesh_box = _make_box_obj("mesh_box", 0.1, 0.1, 0.1) aabb_box = DummyObject( "aabb_only_box", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.05, -0.05, -0.05), max_point=(0.05, 0.05, 0.05)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.05, -0.05, -0.05), max_point=(0.05, 0.05, 0.05)), ) mesh_box.add_relation(On(table)) aabb_box.add_relation(On(table)) diff --git a/isaaclab_arena/tests/test_no_collision_loss.py b/isaaclab_arena/tests/test_no_collision_loss.py index 3b17b5fb33..8132f9dca7 100644 --- a/isaaclab_arena/tests/test_no_collision_loss.py +++ b/isaaclab_arena/tests/test_no_collision_loss.py @@ -5,17 +5,17 @@ """Tests for the RelationSolver built-in no-overlap loss.""" -import math import torch -from isaaclab_arena.relations.loss_primitives import interval_overlap_axis_loss +import pytest + 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 from isaaclab_arena.relations.relation_solver_state import RelationSolverState from isaaclab_arena.relations.relations import IsAnchor, On from isaaclab_arena.tests.dummy_object import DummyObject -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -23,7 +23,7 @@ def _create_box(name: str = "box", size: float = 0.2) -> DummyObject: """Create a small box (local bbox [0,0,0] to [size, size, size]).""" return DummyObject( name=name, - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(size, size, size)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (size, size, size)), ) @@ -31,7 +31,7 @@ def _create_table() -> DummyObject: """Create a table-like object at origin.""" return DummyObject( name="table", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (1.0, 1.0, 0.1)), ) @@ -47,6 +47,44 @@ def _create_no_collision_scene() -> tuple[DummyObject, DummyObject, DummyObject] return table, box_a, box_b +def test_solver_state_rejects_anchor_position_different_from_fixed_pose(): + table, box_a, _ = _create_no_collision_scene() + initial = [{table: (0.1, 0.0, 0.0), box_a: (0.0, 0.0, 0.2)}] + + with pytest.raises(AssertionError, match="must match its fixed Pose position"): + RelationSolverState([table, box_a], initial) + + +@pytest.mark.parametrize( + "rotation", + [ + (0.0, 0.0, 1.0), + (0.0, 0.0, float("nan"), 1.0), + (0.0, 0.0, 0.0, 2.0), + ], +) +def test_solver_state_rejects_invalid_candidate_rotations(rotation): + table, box_a, _ = _create_no_collision_scene() + initial = [{table: (0.0, 0.0, 0.0), box_a: (0.0, 0.0, 0.2)}] + + with pytest.raises(AssertionError, match="Candidate rotations"): + RelationSolverState([table, box_a], initial, rotations=[{box_a: rotation}]) + + +@pytest.mark.parametrize("invalid_key_kind", ["anchor", "unknown"]) +def test_solver_state_rejects_rotation_keys_for_non_optimizable_objects(invalid_key_kind): + table, box_a, _ = _create_no_collision_scene() + initial = [{table: (0.0, 0.0, 0.0), box_a: (0.0, 0.0, 0.2)}] + invalid_key = table if invalid_key_kind == "anchor" else _create_box("unknown") + + with pytest.raises(AssertionError, match="Rotation keys must belong to optimizable objects"): + RelationSolverState( + [table, box_a], + initial, + rotations=[{invalid_key: (0.0, 0.0, 0.0, 1.0)}], + ) + + def test_solver_uses_rotated_bbox_for_collision(): """Test that a yaw-rotated env bbox passed to solve() changes the no-overlap loss (solver consumes it).""" table = _create_table() @@ -57,11 +95,11 @@ def test_solver_uses_rotated_bbox_for_collision(): # the built-in no-overlap loss between the two non-anchors. long_box = DummyObject( name="long_box", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.3, -0.05, -0.05), max_point=(0.3, 0.05, 0.05)), + bounding_box=OrientedBoundingBox.from_min_max((-0.3, -0.05, -0.05), (0.3, 0.05, 0.05)), ) cube = DummyObject( name="cube", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.05, -0.05, -0.05), max_point=(0.05, 0.05, 0.05)), + bounding_box=OrientedBoundingBox.from_min_max((-0.05, -0.05, -0.05), (0.05, 0.05, 0.05)), ) objects = [table, long_box, cube] # Boxes sit high above the table (z=0.5) so neither collides with the table. @@ -74,9 +112,9 @@ def test_solver_uses_rotated_bbox_for_collision(): assert solver.last_loss_per_env is not None loss_unrotated = solver.last_loss_per_env[0].item() - # Hand the solver a 90° conservative bbox for the long box via the env_bboxes channel. - rotated = {long_box: long_box.get_bounding_box().rotated_around_z(math.pi / 2)} - solver.solve(objects, initial, env_bboxes=rotated) + # Hand the solver a 90° candidate rotation through the quaternion channel. + quarter_turn = (0.0, 0.0, 2**-0.5, 2**-0.5) + solver.solve(objects, initial, rotations=[{long_box: quarter_turn}]) assert solver.last_loss_per_env is not None loss_rotated = solver.last_loss_per_env[0].item() @@ -90,30 +128,16 @@ def _single_pair_no_overlap_loss( slope: float, clearance_m: float, child_pos: torch.Tensor, - child_bbox: AxisAlignedBoundingBox, - parent_world_bbox: AxisAlignedBoundingBox, + child_bbox: OrientedBoundingBox, + parent_world_bbox: OrientedBoundingBox, ) -> 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) + child_world = child_bbox.translated(child_pos) + total_loss = slope * child_world.penetration(parent_world_bbox, clearance_m) return total_loss.squeeze(0) if single_input else total_loss @@ -198,7 +222,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 NoCollision loss scales with slope times penetration distance.""" box_a = _create_box("box_a") box_b = _create_box("box_b") @@ -214,15 +238,30 @@ 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_relation_solver_uses_configured_collision_loss_slope(): + """Solver configuration controls the shared OBB and mesh penetration scale.""" + params = RelationSolverParams(collision_loss_slope=25.0, verbose=False) + solver = RelationSolver(params=params) + + assert solver._no_collision_strategy.slope == 25.0 # pyright: ignore[reportPrivateUsage] + + +def test_default_collision_loss_dominates_strongest_relation(): + """Default collision penetration is weighted ten times above an equal On violation.""" + params = RelationSolverParams() + + assert params.collision_loss_slope == 10.0 * params.strategies[On].slope + + +def test_no_collision_loss_minimum_penetration_formula(): + """NoCollision loss equals slope times minimum SAT penetration.""" 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) + # Overlap depth is 0.15 m along every axis. + expected_loss = 10.0 * 0.15 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 @@ -277,7 +316,7 @@ def test_solver_respects_clearance_m(): bbox_b = box_b.get_bounding_box().translated(pos_b) assert not bbox_a.overlaps( - bbox_b, margin=0.05 + bbox_b, clearance_m=0.05 ).item(), f"Boxes should be at least 5 cm apart; box_a at {pos_a}, box_b at {pos_b}" @@ -381,9 +420,9 @@ def test_no_collision_loss_multi_env_shape_and_values(): box_a = _create_box("box_a") child_pos = torch.tensor([[0.0, 0.0, 0.0], [0.1, 0.1, 0.0]]) - parent_world_bbox = AxisAlignedBoundingBox( - min_point=torch.tensor([[1.0, 0.0, 0.0], [0.05, 0.05, 0.0]]), - max_point=torch.tensor([[1.2, 0.2, 0.2], [0.25, 0.25, 0.2]]), + parent_world_bbox = OrientedBoundingBox.from_min_max( + torch.tensor([[1.0, 0.0, 0.0], [0.05, 0.05, 0.0]]), + torch.tensor([[1.2, 0.2, 0.2], [0.25, 0.25, 0.2]]), ) loss = _single_pair_no_overlap_loss( @@ -485,6 +524,28 @@ def test_vectorized_no_overlap_matches_reference_non_anchor_pairs(): _assert_vectorized_matches_reference(objects, initial_positions, expect_positive=True, expected_pair_count=2) +def test_coincident_movable_boxes_receive_opposite_gradients(): + """Directed coincident-box losses separate two movable subjects instead of translating both.""" + table, box_a, box_b = _create_no_collision_scene() + initial_positions = [{ + table: (0.0, 0.0, 0.0), + box_a: (0.3, 0.3, 0.11), + box_b: (0.3, 0.3, 0.11), + }] + solver = RelationSolver(params=RelationSolverParams(verbose=False)) + state = RelationSolverState([table, box_a, box_b], initial_positions, device=torch.device("cpu")) + + loss = solver._compute_no_overlap_loss(state) # pyright: ignore[reportPrivateUsage] + loss.sum().backward() + + assert state.optimizable_positions is not None + gradients = state.optimizable_positions.grad + assert gradients is not None + assert torch.isfinite(gradients).all() + assert torch.count_nonzero(gradients[0, 0]).item() > 0 + torch.testing.assert_close(gradients[0, 0], -gradients[0, 1]) + + def test_vectorized_no_overlap_matches_reference_anchor_pairs(): """Free (non-On) boxes overlapping the anchor exercise the anchor-pair + expand branch.""" table = _create_table() @@ -615,18 +676,14 @@ def test_solver_profile_zero_iters_does_not_raise(): def test_compute_loss_batched_direct(): - """compute_loss_batched returns (num_pairs, batch_size) loss for known world-space extents.""" + """compute_loss_batched scales (P, N) penetration distances.""" strategy = NoCollisionLossStrategy(slope=10.0) - # Two pairs, batch_size=1. Pair 0 overlaps by 0.1 on each axis; pair 1 is fully separated. - subject_min = torch.tensor([[[0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0]]]) - subject_max = torch.tensor([[[0.2, 0.2, 0.2]], [[0.2, 0.2, 0.2]]]) - obstacle_min = torch.tensor([[[0.1, 0.1, 0.1]], [[1.0, 1.0, 1.0]]]) - obstacle_max = torch.tensor([[[0.3, 0.3, 0.3]], [[1.2, 1.2, 1.2]]]) + penetration = torch.tensor([[0.1], [0.0]]) - loss = strategy.compute_loss_batched(0.0, subject_min, subject_max, obstacle_min, obstacle_max) + loss = strategy.compute_loss_batched(penetration) 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(1.0), rtol=1e-4) assert torch.isclose(loss[1, 0], torch.tensor(0.0), atol=1e-6) diff --git a/isaaclab_arena/tests/test_object_mass_variation.py b/isaaclab_arena/tests/test_object_mass_variation.py index 109725c54f..b7462fed52 100644 --- a/isaaclab_arena/tests/test_object_mass_variation.py +++ b/isaaclab_arena/tests/test_object_mass_variation.py @@ -67,7 +67,7 @@ def _test_object_mass_variation_registration(simulation_app): from isaaclab_arena.assets.object_reference import ObjectReference from isaaclab_arena.assets.object_set import RigidObjectSet from isaaclab_arena.assets.registries import AssetRegistry - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose registry = AssetRegistry() @@ -82,8 +82,8 @@ def _test_object_mass_variation_registration(simulation_app): can_a = Object(name="can_a", object_type=ObjectType.RIGID, usd_path="/tmp/can_a.usd") can_b = Object(name="can_b", object_type=ObjectType.RIGID, usd_path="/tmp/can_b.usd") - can_a.bounding_box = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.2)) - can_b.bounding_box = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.3)) + can_a.bounding_box = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.2)) + can_b.bounding_box = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.3)) with ( patch("isaaclab_arena.assets.object_set.detect_object_type", return_value=ObjectType.RIGID), patch("isaaclab_arena.assets.object_set.find_shallowest_rigid_body", return_value="/rigid"), diff --git a/isaaclab_arena/tests/test_object_placer_init.py b/isaaclab_arena/tests/test_object_placer_init.py index 97104e3675..2da34645a1 100644 --- a/isaaclab_arena/tests/test_object_placer_init.py +++ b/isaaclab_arena/tests/test_object_placer_init.py @@ -10,14 +10,14 @@ from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, NextTo, On, Side from isaaclab_arena.tests.dummy_object import DummyObject -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose def _make_desk(): desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (1.0, 1.0, 0.1)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) @@ -33,7 +33,7 @@ def test_on_init_x_y_within_parent_footprint(): desk = _make_desk() box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=0.01)) @@ -45,10 +45,12 @@ def test_on_init_x_y_within_parent_footprint(): child_bbox = box.get_bounding_box() desk_world = desk.get_world_bounding_box() - assert x + child_bbox.min_point[0, 0] >= desk_world.min_point[0, 0] - 1e-6 - assert x + child_bbox.max_point[0, 0] <= desk_world.max_point[0, 0] + 1e-6 - assert y + child_bbox.min_point[0, 1] >= desk_world.min_point[0, 1] - 1e-6 - assert y + child_bbox.max_point[0, 1] <= desk_world.max_point[0, 1] + 1e-6 + child_min, child_max = child_bbox.get_axis_aligned_bounds() + desk_min, desk_max = desk_world.get_axis_aligned_bounds() + assert x + child_min[0, 0] >= desk_min[0, 0] - 1e-6 + assert x + child_max[0, 0] <= desk_max[0, 0] + 1e-6 + assert y + child_min[0, 1] >= desk_min[0, 1] - 1e-6 + assert y + child_max[0, 1] <= desk_max[0, 1] + 1e-6 def test_on_init_z_places_bottom_at_parent_top(): @@ -57,7 +59,7 @@ def test_on_init_z_places_bottom_at_parent_top(): clearance_m = 0.01 box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=clearance_m)) @@ -69,8 +71,10 @@ def test_on_init_z_places_bottom_at_parent_top(): child_bbox = box.get_bounding_box() desk_world = desk.get_world_bounding_box() - child_bottom = z + child_bbox.min_point[0, 2] - expected_bottom = desk_world.max_point[0, 2] + clearance_m + child_min, _ = child_bbox.get_axis_aligned_bounds() + _, desk_max = desk_world.get_axis_aligned_bounds() + child_bottom = z + child_min[0, 2] + expected_bottom = desk_max[0, 2] + clearance_m assert abs(child_bottom - expected_bottom) < 1e-6 @@ -78,13 +82,13 @@ def test_on_init_uses_env_specific_parent_bbox(): """Object with On(anchor set) should initialize against that env's assigned bbox.""" table_set = DummyObject( name="table_set", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 2.0, 0.5)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (2.0, 2.0, 0.5)), ) table_set.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) table_set.add_relation(IsAnchor()) - small_table_bbox = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.1)) - box_bbox = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.05, 0.05, 0.05)) + small_table_bbox = OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.3, 0.3, 0.1)) + box_bbox = OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.05, 0.05, 0.05)) box = DummyObject(name="box", bounding_box=box_bbox) box.add_relation(On(table_set, clearance_m=0.02)) @@ -96,23 +100,25 @@ def test_on_init_uses_env_specific_parent_bbox(): ) x, y, z = positions[box] - assert small_table_bbox.min_point[0, 0] <= x <= small_table_bbox.max_point[0, 0] - assert small_table_bbox.min_point[0, 1] <= y <= small_table_bbox.max_point[0, 1] - assert abs(z - (small_table_bbox.max_point[0, 2] + 0.02 - box_bbox.min_point[0, 2])) < 1e-6 + table_min, table_max = small_table_bbox.get_axis_aligned_bounds() + box_min, _ = box_bbox.get_axis_aligned_bounds() + assert table_min[0, 0] <= x <= table_max[0, 0] + assert table_min[0, 1] <= y <= table_max[0, 1] + assert abs(z - (table_max[0, 2] + 0.02 - box_min[0, 2])) < 1e-6 def test_on_init_clamps_to_center_when_child_wider_than_parent(): """Object wider than its On parent in X/Y is clamped to parent center, not an invalid range.""" desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.1, 0.1, 0.1)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) big_box = DummyObject( name="big_box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.5, 0.5, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.5, 0.5, 0.2)), ) big_box.add_relation(On(desk, clearance_m=0.0)) @@ -122,8 +128,8 @@ def test_on_init_clamps_to_center_when_child_wider_than_parent(): x, y, _ = positions[big_box] desk_world = desk.get_world_bounding_box() - center_x = (desk_world.min_point[0, 0] + desk_world.max_point[0, 0]) / 2.0 - center_y = (desk_world.min_point[0, 1] + desk_world.max_point[0, 1]) / 2.0 + center_x = desk_world.center[0, 0] + center_y = desk_world.center[0, 1] assert abs(x - center_x) < 1e-6 assert abs(y - center_y) < 1e-6 @@ -134,7 +140,7 @@ def test_no_on_relation_initializes_at_anchor_center(): desk = _make_desk() box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.2, 0.2, 0.2)), ) box.add_relation(NextTo(desk, side=Side.POSITIVE_X, distance_m=0.05)) @@ -156,13 +162,13 @@ def test_on_non_anchor_parent_with_anchor_grandparent_uses_proxy(): desk = _make_desk() plate = DummyObject( name="plate", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.02)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.3, 0.3, 0.02)), ) plate.add_relation(On(desk, clearance_m=0.01)) mug = DummyObject( name="mug", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.12)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.1, 0.1, 0.12)), ) mug.add_relation(On(plate, clearance_m=0.0)) @@ -174,11 +180,13 @@ def test_on_non_anchor_parent_with_anchor_grandparent_uses_proxy(): desk_world = desk.get_world_bounding_box() # Mug's parent (plate) is non-anchor with On(desk): uses desk's bbox as proxy - assert desk_world.min_point[0, 0] <= x <= desk_world.max_point[0, 0] - assert desk_world.min_point[0, 1] <= y <= desk_world.max_point[0, 1] + desk_min, desk_max = desk_world.get_axis_aligned_bounds() + assert desk_min[0, 0] <= x <= desk_max[0, 0] + assert desk_min[0, 1] <= y <= desk_max[0, 1] # Z: desk top (0.1) + clearance (0.0) - mug bbox min_z (0.0) = 0.1 mug_bbox = mug.get_bounding_box() - assert abs(z - (desk_world.max_point[0, 2] + 0.0 - mug_bbox.min_point[0, 2])) < 1e-6 + mug_min, _ = mug_bbox.get_axis_aligned_bounds() + assert abs(z - (desk_max[0, 2] - mug_min[0, 2])) < 1e-6 def test_on_non_anchor_parent_without_on_uses_fallback_bbox(): @@ -186,13 +194,13 @@ def test_on_non_anchor_parent_without_on_uses_fallback_bbox(): desk = _make_desk() stand = DummyObject( name="stand", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.5)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.3, 0.3, 0.5)), ) stand.add_relation(NextTo(desk, side=Side.POSITIVE_X, distance_m=0.1)) mug = DummyObject( name="mug", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.12)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.1, 0.1, 0.12)), ) mug.add_relation(On(stand, clearance_m=0.0)) @@ -204,10 +212,12 @@ def test_on_non_anchor_parent_without_on_uses_fallback_bbox(): desk_world = desk.get_world_bounding_box() # Parent (stand) has no On relation: falls back to desk's world bbox - assert desk_world.min_point[0, 0] <= x <= desk_world.max_point[0, 0] - assert desk_world.min_point[0, 1] <= y <= desk_world.max_point[0, 1] + desk_min, desk_max = desk_world.get_axis_aligned_bounds() + assert desk_min[0, 0] <= x <= desk_max[0, 0] + assert desk_min[0, 1] <= y <= desk_max[0, 1] # Z: desk.max_z (fallback) + clearance (0.0) - mug.min_z (0.0) = 0.1 - assert abs(z - (desk_world.max_point[0, 2] + 0.0 - mug.get_bounding_box().min_point[0, 2])) < 1e-6 + mug_min, _ = mug.get_bounding_box().get_axis_aligned_bounds() + assert abs(z - (desk_max[0, 2] - mug_min[0, 2])) < 1e-6 def test_on_init_reproducible_with_placement_seed(): @@ -219,7 +229,7 @@ def _run(): desk = _make_desk() box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=0.01)) (result,) = ObjectPlacer(params=params).place([desk, box]) diff --git a/isaaclab_arena/tests/test_object_placer_reproducibility.py b/isaaclab_arena/tests/test_object_placer_reproducibility.py index 08eb807283..0c4f99962e 100644 --- a/isaaclab_arena/tests/test_object_placer_reproducibility.py +++ b/isaaclab_arena/tests/test_object_placer_reproducibility.py @@ -17,7 +17,7 @@ from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, NextTo, On, RotateAroundSolution, Side from isaaclab_arena.tests.dummy_object import DummyObject -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, get_random_pose_within_bounding_box +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox, get_random_pose_within_bounding_box from isaaclab_arena.utils.pose import Pose, PosePerEnv from isaaclab_arena.utils.yaw import rotate_quat_by_yaw, wrap_angle_to_pi @@ -27,18 +27,18 @@ def _create_test_objects(): desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (1.0, 1.0, 0.1)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) box1 = DummyObject( name="box1", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.2, 0.2, 0.2)), ) box2 = DummyObject( name="box2", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.15)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.15, 0.15, 0.15)), ) box1.add_relation(On(desk, clearance_m=0.01)) @@ -51,7 +51,7 @@ def _create_test_objects(): def test_get_random_pose_same_seed_produces_identical_result(): """Test that get_random_pose_within_bounding_box with same seed produces identical poses.""" - bbox = AxisAlignedBoundingBox(min_point=(-1.0, -1.0, 0.0), max_point=(1.0, 1.0, 1.0)) + bbox = OrientedBoundingBox.from_min_max((-1.0, -1.0, 0.0), (1.0, 1.0, 1.0)) pose1 = get_random_pose_within_bounding_box(bbox, seed=42) pose2 = get_random_pose_within_bounding_box(bbox, seed=42) @@ -303,7 +303,7 @@ def test_random_yaw_init_applied_yaw_matches_selected_candidate(): (result,) = placer.place([desk, box1, box2], num_envs=1) for box in (box1, box2): applied = _yaw_rad_from_quat(box.get_initial_pose().rotation_xyzw) - assert abs(wrap_angle_to_pi(applied - result.orientations[box])) < 1e-5 + assert abs(wrap_angle_to_pi(applied - _yaw_rad_from_quat(result.rotations[box]))) < 1e-5 def test_random_yaw_init_composes_marker_yaw(): @@ -316,32 +316,27 @@ def test_random_yaw_init_composes_marker_yaw(): ) (result,) = placer.place([desk, box1, box2], num_envs=1) applied = _yaw_rad_from_quat(box1.get_initial_pose().rotation_xyzw) - # result.orientations now carries total yaw = marker + sampled - assert abs(wrap_angle_to_pi(applied - result.orientations[box1])) < 1e-5 + assert abs(wrap_angle_to_pi(applied - _yaw_rad_from_quat(result.rotations[box1]))) < 1e-5 -def test_roll_pitch_marker_applied_verbatim_without_random_yaw(): - """A roll/pitch marker is excluded from random-yaw sampling and applied verbatim. - - A Z-rotated footprint can't enclose a tilted box, so the object keeps its requested rotation - and receives no sampled yaw even when random_yaw_init is on. - """ - pitch_rad = math.pi / 2 +def test_roll_pitch_marker_survives_place_with_forced_world_yaw(monkeypatch): + """A roll/pitch marker and world yaw reach the result and applied Pose unchanged.""" solver_params = RelationSolverParams(max_iters=5, verbose=False) desk, box1, box2 = _create_test_objects() - box1.add_relation(RotateAroundSolution(pitch_rad=pitch_rad)) + box1.add_relation(RotateAroundSolution(roll_rad=0.4, pitch_rad=-0.3)) placer = ObjectPlacer( params=ObjectPlacerParams(placement_seed=1, solver_params=solver_params, random_yaw_init=True) ) - orientations = placer._generate_initial_orientations([desk, box1, box2], {desk}) - assert box1 not in orientations, "roll/pitch marker object must not receive a sampled yaw" + monkeypatch.setattr("isaaclab_arena.relations.object_placer.get_random_rotation", lambda generator: math.pi / 2) + expected = placer._generate_initial_rotations([desk, box1, box2], {desk})[box1] - placer.place([desk, box1, box2], num_envs=1) - applied = box1.get_initial_pose().rotation_xyzw - expected = RotateAroundSolution(pitch_rad=pitch_rad).get_rotation_xyzw() - assert all( - abs(a - e) < 1e-5 for a, e in zip(applied, expected) - ), f"marker pitch must be applied verbatim; expected {expected}, got {applied}" + (result,) = placer.place([desk, box1, box2]) + applied = box1.get_initial_pose() + + assert isinstance(applied, Pose) + assert abs(expected[0]) > 1e-3 and abs(expected[1]) > 1e-3 + assert result.rotations[box1] == pytest.approx(expected) + assert applied.rotation_xyzw == pytest.approx(expected) def test_marker_yaw_applied_without_random_yaw_init(): @@ -353,8 +348,8 @@ def test_marker_yaw_applied_without_random_yaw_init(): placer = ObjectPlacer( params=ObjectPlacerParams(placement_seed=1, solver_params=solver_params, random_yaw_init=False) ) - orientations = placer._generate_initial_orientations([desk, box1, box2], {desk}) - assert abs(wrap_angle_to_pi(orientations[box1] - marker_yaw)) < 1e-5 + rotations = placer._generate_initial_rotations([desk, box1, box2], {desk}) + assert abs(wrap_angle_to_pi(_yaw_rad_from_quat(rotations[box1]) - marker_yaw)) < 1e-5 placer.place([desk, box1, box2], num_envs=1) applied = _yaw_rad_from_quat(box1.get_initial_pose().rotation_xyzw) assert abs(wrap_angle_to_pi(applied - marker_yaw)) < 1e-5, f"Marker yaw {marker_yaw} must be applied; got {applied}" diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py index 0062bf4066..f4bd4c870e 100644 --- a/isaaclab_arena/tests/test_object_set.py +++ b/isaaclab_arena/tests/test_object_set.py @@ -20,12 +20,12 @@ def _make_object_set_variants(): from isaaclab_arena.assets.object import Object from isaaclab_arena.assets.object_base import ObjectType - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox can_a = Object(name="can_a", object_type=ObjectType.RIGID, usd_path="/tmp/can_a.usd") can_b = Object(name="can_b", object_type=ObjectType.RIGID, usd_path="/tmp/can_b.usd") - bbox_a = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.2)) - bbox_b = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.3)) + bbox_a = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.2)) + bbox_b = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.3)) can_a.bounding_box = bbox_a can_b.bounding_box = bbox_b return can_a, can_b, bbox_a, bbox_b @@ -62,8 +62,8 @@ def _test_object_set_samples_and_stores_variant_indices(simulation_app): assert contact_sensor_cfg.prim_path == f"{obj_set.prim_path}/rigid" per_env_bbox = obj_set.get_bounding_box_per_env(num_envs=4) - assert torch.allclose(per_env_bbox.max_point[0], bbox_b.max_point[0]) - assert torch.allclose(per_env_bbox.max_point[1], bbox_a.max_point[0]) + assert torch.allclose(per_env_bbox.center[0], bbox_b.center[0]) + assert torch.allclose(per_env_bbox.half_extents[1], bbox_a.half_extents[0]) return True @@ -89,8 +89,8 @@ def _test_object_set_default_variant_indices_follow_member_order(simulation_app) assert getattr(spawn_cfg, "random_choice") is False per_env_bbox = obj_set.get_bounding_box_per_env(num_envs=5) - assert torch.allclose(per_env_bbox.max_point[0], bbox_a.max_point[0]) - assert torch.allclose(per_env_bbox.max_point[1], bbox_b.max_point[0]) + assert torch.allclose(per_env_bbox.center[0], bbox_a.center[0]) + assert torch.allclose(per_env_bbox.half_extents[1], bbox_b.half_extents[0]) return True diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index 29b7dfbeda..a454b396e6 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -27,23 +27,23 @@ def _create_test_objects(): from isaaclab_arena.relations.relations import IsAnchor, NextTo, On, Side from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) box1 = DummyObject( name="box1", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) box2 = DummyObject( name="box2", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.15)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.15)), ) box1.add_relation(On(desk, clearance_m=0.01)) @@ -194,12 +194,12 @@ def test_solve_and_place_objects_uses_runtime_pool(): from isaaclab_arena.relations.placement_events import solve_and_place_objects from isaaclab_arena.relations.placement_result import PlacementResult from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox desk, _, _ = _create_test_objects() robot = DummyEmbodiment( name="droid", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.2, -0.2, 0.0), max_point=(0.2, 0.2, 1.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.2, -0.2, 0.0), max_point=(0.2, 0.2, 1.0)), scene_name="robot", ) env = _make_mock_env(num_envs=1) @@ -672,19 +672,19 @@ def test_env_indexed_static_poses_apply_per_env_positions(): from isaaclab_arena.relations.placement_result import PlacementResult from isaaclab_arena.relations.relations import IsAnchor, On from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose, PosePerEnv desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=0.01)) @@ -726,12 +726,12 @@ def test_pooled_placer_falls_back_when_no_valid_layouts(capsys): from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.01, 0.01, 0.01)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.01, 0.01, 0.01)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) @@ -739,11 +739,11 @@ def test_pooled_placer_falls_back_when_no_valid_layouts(capsys): # Two large boxes that cannot both fit On a tiny desk big1 = DummyObject( name="big1", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(5.0, 5.0, 5.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(5.0, 5.0, 5.0)), ) big2 = DummyObject( name="big2", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(5.0, 5.0, 5.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(5.0, 5.0, 5.0)), ) big1.add_relation(On(desk)) big2.add_relation(On(desk)) @@ -768,19 +768,19 @@ def test_pooled_placer_only_falls_back_on_final_batch(capsys): from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.01, 0.01, 0.01)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.01, 0.01, 0.01)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) big = DummyObject( name="big", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(5.0, 5.0, 5.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(5.0, 5.0, 5.0)), ) big.add_relation(On(desk)) @@ -804,23 +804,23 @@ def test_pooled_placer_can_reject_best_loss_fallbacks(): from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.01, 0.01, 0.01)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.01, 0.01, 0.01)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) big1 = DummyObject( name="big1", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(5.0, 5.0, 5.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(5.0, 5.0, 5.0)), ) big2 = DummyObject( name="big2", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(5.0, 5.0, 5.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(5.0, 5.0, 5.0)), ) big1.add_relation(On(desk)) big2.add_relation(On(desk)) @@ -868,7 +868,7 @@ def validate_batch(self, positions, orientations, bboxes, collision_objects): positions=positions[i], final_loss=0.0, attempts=0, - orientations=orientations[i], + rotations=orientations[i], ) for i in range(len(positions)) ] diff --git a/isaaclab_arena/tests/test_pose.py b/isaaclab_arena/tests/test_pose.py index 6c8cccfcf3..b406e4a3ba 100644 --- a/isaaclab_arena/tests/test_pose.py +++ b/isaaclab_arena/tests/test_pose.py @@ -6,6 +6,8 @@ import math import torch +import pytest + from isaaclab_arena.utils.pose import Pose, PosePerEnv from isaaclab_arena.utils.yaw import ( rotate_points_by_yaw, @@ -37,6 +39,14 @@ def test_rotate_quat_by_yaw_composes_and_wraps(): assert rotate_quat_by_yaw(base, 2.0 * math.pi) == base +def test_rotate_quat_by_yaw_precomposes_world_yaw_over_pitch(): + """World yaw multiplies on the left when the base rotation contains pitch.""" + half_sqrt = 2**-0.5 + result = rotate_quat_by_yaw((0.0, half_sqrt, 0.0, half_sqrt), math.pi / 2) + + assert result == pytest.approx((-0.5, 0.5, 0.5, 0.5)) + + def test_pose_composition(): T_B_A = Pose(position_xyz=(1.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) T_C_B = Pose(position_xyz=(2.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) diff --git a/isaaclab_arena/tests/test_position_limits.py b/isaaclab_arena/tests/test_position_limits.py index bc398a3c3d..84a6af79f2 100644 --- a/isaaclab_arena/tests/test_position_limits.py +++ b/isaaclab_arena/tests/test_position_limits.py @@ -11,10 +11,10 @@ from isaaclab_arena.relations.relation_loss_strategies import PositionLimitsLossStrategy from isaaclab_arena.relations.relations import PositionLimits -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox # Dummy bounding box used for all strategy tests (child object is a 0.1m cube at origin) -_DUMMY_BBOX = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) +_DUMMY_BBOX = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)) # ============================================================================= @@ -167,14 +167,14 @@ def test_solver_respects_position_limits(): """Solver moves an object inside the PositionLimits region.""" table = DummyObject( name="table", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 2.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 2.0, 0.1)), ) table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) table.add_relation(IsAnchor()) box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)), ) box.add_relation(On(table, clearance_m=0.01)) box.add_relation(PositionLimits(x_min=0.2, x_max=0.5, y_min=0.2, y_max=0.5)) diff --git a/isaaclab_arena/tests/test_reference_objects.py b/isaaclab_arena/tests/test_reference_objects.py index c99f3c7930..96e4115616 100644 --- a/isaaclab_arena/tests/test_reference_objects.py +++ b/isaaclab_arena/tests/test_reference_objects.py @@ -10,8 +10,10 @@ import traceback from types import SimpleNamespace +import pytest + from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose NUM_STEPS = 50 @@ -40,7 +42,7 @@ def __init__(self): return ObjectReferenceTestKitchenBackground() -def _object_reference_with_cached_bbox(parent_pose: Pose | None, relative_pose: Pose, bbox: AxisAlignedBoundingBox): +def _object_reference_with_cached_bbox(parent_pose: Pose | None, relative_pose: Pose, bbox: OrientedBoundingBox): """Construct an ObjectReference around cached geometry without opening a USD.""" from isaaclab_arena.assets.object_reference import ObjectReference @@ -52,18 +54,205 @@ def _object_reference_with_cached_bbox(parent_pose: Pose | None, relative_pose: def test_object_reference_world_bbox_applies_parent_yaw(): - """Parent yaw, not the prim's relative yaw, rotates the already-local referenced bbox.""" + """The composed reference pose rotates the local bounding box.""" yaw_90 = (0.0, 0.0, 2**-0.5, 2**-0.5) obj_ref = _object_reference_with_cached_bbox( parent_pose=Pose(position_xyz=(10.0, 0.0, 0.0), rotation_xyzw=yaw_90), relative_pose=Pose(position_xyz=(1.0, 2.0, 0.0), rotation_xyzw=yaw_90), - bbox=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.1, 0.05)), + bbox=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.1, 0.05)), ) world_bbox = obj_ref.get_world_bounding_box() + min_point, max_point = world_bbox.get_axis_aligned_bounds() + + assert torch.allclose(min_point, torch.tensor([[7.8, 0.9, 0.0]]), atol=1e-6) + assert torch.allclose(max_point, torch.tensor([[8.0, 1.0, 0.05]]), atol=1e-6) + + +@pytest.mark.parametrize("scale", [(1.0, 0.0, 1.0), (1.0, -1.0, 1.0)]) +def test_object_reference_rejects_non_positive_parent_scale(scale): + """Object references require positive parent scale components.""" + from isaaclab_arena.assets.object_reference import ObjectReference + + parent = SimpleNamespace(scale=scale) + with pytest.raises(AssertionError, match="parent scale must be positive"): + ObjectReference(parent_asset=parent, name="reference") + + +def _test_object_reference_nonuniform_parent_scale_with_rotation(simulation_app) -> bool: + """Reference-local geometry applies R^-1 S R before its rigid world pose.""" + import math + from contextlib import nullcontext + from unittest.mock import patch + + from pxr import Gf, Usd, UsdGeom + + from isaaclab_arena.assets.object_reference import ObjectReference + + raw_vertices = np.array([ + [-1.0, -2.0, -0.5], + [3.0, -2.0, -0.5], + [3.0, 4.0, -0.5], + [-1.0, 4.0, -0.5], + [-1.0, -2.0, 1.5], + [3.0, -2.0, 1.5], + [3.0, 4.0, 1.5], + [-1.0, 4.0, 1.5], + ]) + stage = Usd.Stage.CreateInMemory() + root = UsdGeom.Xform.Define(stage, "/Root") + stage.SetDefaultPrim(root.GetPrim()) + reference = UsdGeom.Mesh.Define(stage, "/Root/Reference") + reference.GetPointsAttr().Set([Gf.Vec3f(*vertex) for vertex in raw_vertices]) + reference.GetFaceVertexCountsAttr().Set([4, 4, 4, 4, 4, 4]) + reference.GetFaceVertexIndicesAttr().Set([ + 0, + 1, + 2, + 3, + 4, + 7, + 6, + 5, + 0, + 4, + 5, + 1, + 1, + 5, + 6, + 2, + 2, + 6, + 7, + 3, + 4, + 0, + 3, + 7, + ]) + UsdGeom.Xformable(reference).AddRotateZOp().Set(90.0) + + scale = np.array([2.0, 5.0, 3.0]) + rotation = np.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) + raw_translation = np.array([1.0, 2.0, 3.0]) + scaled_translation = scale * raw_translation + expected_local = ((raw_vertices @ rotation.T) * scale) @ rotation + + obj_ref = ObjectReference.__new__(ObjectReference) + obj_ref.name = "reference" + obj_ref.parent_asset = SimpleNamespace(usd_path="/tmp/reference.usd", name="parent", initial_pose=None) + obj_ref.prim_path = "{ENV_REGEX_NS}/parent/Reference" + obj_ref._parent_scale = tuple(scale) + obj_ref.initial_pose_relative_to_parent = Pose( + position_xyz=tuple(scaled_translation), + rotation_xyzw=(0.0, 0.0, math.sqrt(0.5), math.sqrt(0.5)), + ) + obj_ref._bounding_box = None + obj_ref._collision_mesh = None + obj_ref._collision_mesh_loaded = False + + with ( + patch("isaaclab_arena.assets.object_reference.open_stage", return_value=nullcontext(stage)), + patch.object( + ObjectReference, + "isaaclab_prim_path_to_original_prim_path", + staticmethod(lambda prim_path, parent, opened_stage: "/Root/Reference"), + ), + ): + local_box = obj_ref.get_bounding_box() + mesh = obj_ref.get_collision_mesh() + + np.testing.assert_allclose(local_box.center.numpy(), [[5.0, 2.0, 1.5]], atol=1e-6) + np.testing.assert_allclose(local_box.half_extents.numpy(), [[10.0, 6.0, 3.0]], atol=1e-6) + assert mesh is not None + np.testing.assert_allclose(mesh.vertices, expected_local, atol=1e-6) + + expected_world = raw_vertices @ rotation.T * scale + scaled_translation + mesh_world = mesh.vertices @ rotation.T + scaled_translation + np.testing.assert_allclose(mesh_world, expected_world, atol=1e-6) + + world_box = obj_ref.get_world_bounding_box() + np.testing.assert_allclose(world_box.center.numpy(), [[0.0, 15.0, 10.5]], atol=1e-6) + np.testing.assert_allclose( + world_box.get_axis_aligned_bounds()[0].numpy(), + [[-6.0, 5.0, 7.5]], + atol=1e-6, + ) + np.testing.assert_allclose( + world_box.get_axis_aligned_bounds()[1].numpy(), + [[6.0, 25.0, 13.5]], + atol=1e-6, + ) + return True + + +def test_object_reference_nonuniform_parent_scale_with_rotation(): + assert run_simulation_app_function( + _test_object_reference_nonuniform_parent_scale_with_rotation, + headless=HEADLESS, + ) + + +def _test_rotated_reference_local_bbox_is_not_double_rotated(simulation_app) -> bool: + """A rotated reference keeps axis-aligned geometry in its own frame.""" + import math + + from pxr import Gf, Usd, UsdGeom + + from isaaclab_arena.assets.object_reference import ObjectReference + from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_prim + + stage = Usd.Stage.CreateInMemory() + root = UsdGeom.Xform.Define(stage, "/Root") + stage.SetDefaultPrim(root.GetPrim()) + reference = UsdGeom.Mesh.Define(stage, "/Root/Reference") + reference.GetPointsAttr().Set([ + Gf.Vec3f(x, y, z) + for x, y, z in ( + (-2.0, -1.0, -0.5), + (2.0, -1.0, -0.5), + (2.0, 1.0, -0.5), + (-2.0, 1.0, -0.5), + (-2.0, -1.0, 0.5), + (2.0, -1.0, 0.5), + (2.0, 1.0, 0.5), + (-2.0, 1.0, 0.5), + ) + ]) + UsdGeom.Xformable(reference).AddRotateZOp().Set(45.0) + + local_bbox = compute_local_bounding_box_from_prim(stage, "/Root/Reference") + torch.testing.assert_close(local_bbox.half_extents, torch.tensor([[2.0, 1.0, 0.5]])) + + half_angle = math.pi / 8.0 + obj_ref = ObjectReference.__new__(ObjectReference) + obj_ref.parent_asset = SimpleNamespace(initial_pose=None) + obj_ref.initial_pose_relative_to_parent = Pose( + position_xyz=(0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, math.sin(half_angle), math.cos(half_angle)), + ) + obj_ref._bounding_box = local_bbox + + minimum, maximum = obj_ref.get_world_bounding_box().get_axis_aligned_bounds() + xy_half_extent = 3.0 / math.sqrt(2.0) + torch.testing.assert_close( + minimum, + torch.tensor([[-xy_half_extent, -xy_half_extent, -0.5]]), + atol=1e-6, + rtol=0, + ) + torch.testing.assert_close( + maximum, + torch.tensor([[xy_half_extent, xy_half_extent, 0.5]]), + atol=1e-6, + rtol=0, + ) + return True + - assert torch.allclose(world_bbox.min_point, torch.tensor([[7.9, 1.0, 0.0]]), atol=1e-6) - assert torch.allclose(world_bbox.max_point, torch.tensor([[8.0, 1.2, 0.05]]), atol=1e-6) +def test_rotated_reference_local_bbox_is_not_double_rotated(): + assert run_simulation_app_function(_test_rotated_reference_local_bbox_is_not_double_rotated, headless=HEADLESS) def test_object_reference_get_collision_mesh_extracts_referenced_prim(monkeypatch): @@ -77,7 +266,8 @@ def test_object_reference_get_collision_mesh_extracts_referenced_prim(monkeypatc obj_ref = ObjectReference.__new__(ObjectReference) obj_ref.parent_asset = SimpleNamespace(usd_path="/tmp/kitchen.usd", name="kitchen") obj_ref.prim_path = "{ENV_REGEX_NS}/kitchen/counter" - obj_ref._parent_scale = (2.0, 1.0, 1.0) + obj_ref._parent_scale = (2.0, 2.0, 2.0) + obj_ref.initial_pose_relative_to_parent = Pose.identity() obj_ref._collision_mesh = None obj_ref._collision_mesh_loaded = False @@ -112,7 +302,7 @@ def fake_extract(stage, prim_path, scale): assert obj_ref.get_collision_mesh() is expected_mesh assert calls == { "opened": "/tmp/kitchen.usd", - "extract": ("/World/counter", (2.0, 1.0, 1.0)), + "extract": ("/World/counter", (1.0, 1.0, 1.0)), } @@ -248,13 +438,14 @@ def test_object_reference_world_bbox_without_parent_pose_uses_reference_pose(): obj_ref = _object_reference_with_cached_bbox( parent_pose=None, relative_pose=Pose(position_xyz=(1.0, 2.0, 3.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)), - bbox=AxisAlignedBoundingBox(min_point=(-0.1, -0.2, 0.0), max_point=(0.1, 0.2, 0.3)), + bbox=OrientedBoundingBox.from_min_max(min_point=(-0.1, -0.2, 0.0), max_point=(0.1, 0.2, 0.3)), ) world_bbox = obj_ref.get_world_bounding_box() + min_point, max_point = world_bbox.get_axis_aligned_bounds() - assert torch.allclose(world_bbox.min_point, torch.tensor([[0.9, 1.8, 3.0]]), atol=1e-6) - assert torch.allclose(world_bbox.max_point, torch.tensor([[1.1, 2.2, 3.3]]), atol=1e-6) + assert torch.allclose(min_point, torch.tensor([[0.9, 1.8, 3.0]]), atol=1e-6) + assert torch.allclose(max_point, torch.tensor([[1.1, 2.2, 3.3]]), atol=1e-6) def get_test_scene(): diff --git a/isaaclab_arena/tests/test_relation_loss_strategies.py b/isaaclab_arena/tests/test_relation_loss_strategies.py index b2558cc628..5653461a47 100644 --- a/isaaclab_arena/tests/test_relation_loss_strategies.py +++ b/isaaclab_arena/tests/test_relation_loss_strategies.py @@ -14,7 +14,7 @@ from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, NextTo, NotNextTo, On, Side from isaaclab_arena.tests.dummy_object import DummyObject -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -23,7 +23,7 @@ def _create_table(): return DummyObject( name="table", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (1.0, 1.0, 0.1)), ) @@ -32,7 +32,7 @@ def _create_box(): return DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (0.2, 0.2, 0.2)), ) @@ -170,7 +170,7 @@ def test_on_loss_strategy_oversized_child_keeps_plateau_without_margin(): table = _create_table() # X extent [0, 1] wide_box = DummyObject( name="wide_box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.4, 0.2, 0.15)), + bounding_box=OrientedBoundingBox.from_min_max((0.0, 0.0, 0.0), (1.4, 0.2, 0.15)), ) strategy = OnLossStrategy(slope=10.0) relation = On(table, clearance_m=0.01, edge_margin_m=0.0) @@ -310,9 +310,9 @@ def test_on_loss_strategy_multi_env_shape_and_values(): strategy = OnLossStrategy(slope=10.0) child_pos = torch.tensor([[0.4, 0.4, 0.11], [0.4, 0.4, 0.5]]) - parent_world_bbox = AxisAlignedBoundingBox( - min_point=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), - max_point=torch.tensor([[1.0, 1.0, 0.1], [1.0, 1.0, 0.1]]), + parent_world_bbox = OrientedBoundingBox.from_min_max( + torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), + torch.tensor([[1.0, 1.0, 0.1], [1.0, 1.0, 0.1]]), ) loss = strategy.compute_loss(relation, child_pos, box.bounding_box, parent_world_bbox) diff --git a/isaaclab_arena/tests/test_relation_solver_background_collision.py b/isaaclab_arena/tests/test_relation_solver_background_collision.py index 83dfb443ac..1eee09ca97 100644 --- a/isaaclab_arena/tests/test_relation_solver_background_collision.py +++ b/isaaclab_arena/tests/test_relation_solver_background_collision.py @@ -18,12 +18,12 @@ def _make_desk(): """Anchor desk, 2m x 1m wide so a box has room to relocate along X away from the obstacle.""" from isaaclab_arena.relations.relations import IsAnchor from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 1.0, 0.1)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) @@ -33,23 +33,23 @@ def _make_desk(): def _make_box(name: str = "box"): """A 0.3m cube to place On the desk (smaller than each desk half so a valid spot exists).""" from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox return DummyObject( name=name, - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3)), ) def _make_background(): """Return an obstacle narrower than the desk to preserve a non-zero escape gradient.""" from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose background = DummyObject( name="cabinet", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.5, 1.0, 1.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.5, 1.0, 1.0)), ) background.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) return background @@ -60,13 +60,13 @@ def _mesh_box(name: str, extents: tuple[float, float, float], position: tuple[fl import trimesh from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose half = tuple(e / 2.0 for e in extents) obj = DummyObject( name=name, - bounding_box=AxisAlignedBoundingBox( + bounding_box=OrientedBoundingBox.from_min_max( min_point=(-half[0], -half[1], -half[2]), max_point=(half[0], half[1], half[2]), ), @@ -179,8 +179,9 @@ def test_background_collision_objects_treat_background_none_pose_as_identity(mon collision_objects = make_fixed_collision_objects([kitchen]) assert len(collision_objects) == 1 - assert torch.allclose(collision_objects[0].get_bounding_box().min_point, torch.tensor([[-0.1, -0.1, -0.1]])) - assert torch.allclose(collision_objects[0].get_bounding_box().max_point, torch.tensor([[0.1, 0.1, 0.1]])) + min_point, max_point = collision_objects[0].get_bounding_box().get_axis_aligned_bounds() + assert torch.allclose(min_point, torch.tensor([[-0.1, -0.1, -0.1]])) + assert torch.allclose(max_point, torch.tensor([[0.1, 0.1, 0.1]])) def test_background_collision_objects_reject_bbox_whole_background(): @@ -292,7 +293,7 @@ def test_solver_pushes_object_off_background_obstacle(): result = solver.solve([desk, box], initial_positions, collision_objects=[background])[0] box_world = box.get_bounding_box().translated(result[box]) - assert not box_world.overlaps(background.get_world_bounding_box(), margin=solver_params.clearance_m).item() + assert not box_world.overlaps(background.get_world_bounding_box(), clearance_m=solver_params.clearance_m).item() def test_solve_without_collision_objects_is_a_noop(): @@ -309,7 +310,7 @@ def test_solve_without_collision_objects_is_a_noop(): solver = RelationSolver(RelationSolverParams(verbose=False, save_position_history=False)) result = solver.solve([desk, box], initial_positions)[0] - box_bottom_z = box.get_bounding_box().min_point[0, 2].item() + result[box][2] + box_bottom_z = box.get_bounding_box().get_axis_aligned_bounds()[0][0, 2].item() + result[box][2] assert abs(box_bottom_z - 0.1) < 0.05 @@ -342,10 +343,10 @@ def test_validate_no_overlap_rejects_background_overlap(): env_bboxes = {desk: desk.get_bounding_box(), box: box.get_bounding_box()} overlapping = {desk: (0.0, 0.0, 0.0), box: (0.3, 0.3, 0.1)} - assert not validator._validate_no_overlap(overlapping, env_bboxes, [background]) + assert not validator._validate_no_overlap(overlapping, env_bboxes, collision_objects=[background]) clear = {desk: (0.0, 0.0, 0.0), box: (1.5, 0.3, 0.1)} - assert validator._validate_no_overlap(clear, env_bboxes, [background]) + assert validator._validate_no_overlap(clear, env_bboxes, collision_objects=[background]) def _test_get_passive_collision_objects_filters(simulation_app) -> bool: @@ -454,7 +455,7 @@ def test_object_placer_place_forwards_collision_objects(): assert result.success box_world = box.get_bounding_box().translated(result.positions[box]) - assert not box_world.overlaps(background.get_world_bounding_box(), margin=solver_params.clearance_m).item() + assert not box_world.overlaps(background.get_world_bounding_box(), clearance_m=solver_params.clearance_m).item() def test_pooled_object_placer_forwards_collision_objects(): @@ -485,7 +486,7 @@ def test_pooled_object_placer_forwards_collision_objects(): layout = pool.sample_with_replacement(1)[0] box_world = box.get_bounding_box().translated(layout.positions[box]) - assert not box_world.overlaps(background.get_world_bounding_box(), margin=solver_params.clearance_m).item() + assert not box_world.overlaps(background.get_world_bounding_box(), clearance_m=solver_params.clearance_m).item() def test_pooled_object_placer_multi_env_avoids_obstacle(): @@ -520,7 +521,7 @@ def test_pooled_object_placer_multi_env_avoids_obstacle(): assert len(layouts) == num_envs for layout in layouts: box_world = box.get_bounding_box().translated(layout.positions[box]) - assert not box_world.overlaps(background.get_world_bounding_box(), margin=solver_params.clearance_m).item() + assert not box_world.overlaps(background.get_world_bounding_box(), clearance_m=solver_params.clearance_m).item() def test_arena_env_builder_forwards_background_collisions_by_default(monkeypatch): @@ -652,11 +653,11 @@ def test_relation_placement_includes_background_mesh_for_object_mesh_override(mo 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 + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox mesh_object = DummyObject( "mesh_object", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.1, -0.1, -0.1), max_point=(0.1, 0.1, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(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 @@ -695,13 +696,13 @@ def test_relation_placement_includes_background_mesh_for_background_override(mon 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 + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox background = Background.__new__(Background) background.collision_mode = CollisionMode.MESH placed_object = DummyObject( "placed_object", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.1, -0.1, -0.1), max_point=(0.1, 0.1, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.1, -0.1, -0.1), max_point=(0.1, 0.1, 0.1)), ) placed_object.add_relation(IsAnchor()) calls = {} @@ -741,13 +742,13 @@ def test_relation_placement_skips_background_mesh_for_default_bbox(monkeypatch): 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 + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox background = Background.__new__(Background) background.collision_mode = None placed_object = DummyObject( "placed_object", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.1, -0.1, -0.1), max_point=(0.1, 0.1, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.1, -0.1, -0.1), max_point=(0.1, 0.1, 0.1)), ) placed_object.add_relation(IsAnchor()) calls = {} diff --git a/isaaclab_arena/tests/test_relation_solver_embodiment.py b/isaaclab_arena/tests/test_relation_solver_embodiment.py index 9fc2484097..ada0170038 100644 --- a/isaaclab_arena/tests/test_relation_solver_embodiment.py +++ b/isaaclab_arena/tests/test_relation_solver_embodiment.py @@ -12,14 +12,14 @@ from isaaclab_arena.relations.relations import IsAnchor, On from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment from isaaclab_arena.tests.dummy_object import DummyObject -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose, PosePerEnv def _make_floor_and_robot(): floor = DummyObject( name="floor", - bounding_box=AxisAlignedBoundingBox( + bounding_box=OrientedBoundingBox.from_min_max( min_point=(-2.0, -2.0, -0.05), max_point=(2.0, 2.0, 0.0), ), @@ -28,7 +28,7 @@ def _make_floor_and_robot(): floor.add_relation(IsAnchor()) robot = DummyEmbodiment( name="robot", - bounding_box=AxisAlignedBoundingBox( + bounding_box=OrientedBoundingBox.from_min_max( min_point=(-0.2, -0.2, 0.0), max_point=(0.2, 0.2, 1.2), ), @@ -61,7 +61,7 @@ def test_batched_embodiment_placement_stores_per_env_poses(): def test_world_bounding_box_applies_positive_quarter_turn(): asset = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 1.0, 1.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 1.0, 1.0)), initial_pose=Pose( position_xyz=(3.0, 4.0, 0.0), rotation_xyzw=(0.0, 0.0, 2**-0.5, 2**-0.5), @@ -69,6 +69,7 @@ def test_world_bounding_box_applies_positive_quarter_turn(): ) world_bbox = asset.get_world_bounding_box() + min_point, max_point = world_bbox.get_axis_aligned_bounds() - assert torch.allclose(world_bbox.min_point, torch.tensor([[2.0, 4.0, 0.0]])) - assert torch.allclose(world_bbox.max_point, torch.tensor([[3.0, 6.0, 1.0]])) + assert torch.allclose(min_point, torch.tensor([[2.0, 4.0, 0.0]])) + assert torch.allclose(max_point, torch.tensor([[3.0, 6.0, 1.0]])) diff --git a/isaaclab_arena/tests/test_relation_solver_interface.py b/isaaclab_arena/tests/test_relation_solver_interface.py index 43d32a534b..0074b78093 100644 --- a/isaaclab_arena/tests/test_relation_solver_interface.py +++ b/isaaclab_arena/tests/test_relation_solver_interface.py @@ -8,18 +8,18 @@ import pytest from isaaclab_arena.tests.dummy_object import DummyObject -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox def _make_desk(): from isaaclab_arena.relations.relations import IsAnchor from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) @@ -28,11 +28,11 @@ def _make_desk(): def _make_box(name: str = "box"): from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox return DummyObject( name=name, - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) @@ -75,12 +75,12 @@ def test_solve_and_apply_relation_placement_requires_unique_asset_names(): def test_solve_and_apply_relation_placement_rejects_scene_name_collision(): from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox embodiment = DummyEmbodiment( name="droid", scene_name="robot", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.2, -0.2, 0.0), max_point=(0.2, 0.2, 1.0)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(-0.2, -0.2, 0.0), max_point=(0.2, 0.2, 1.0)), ) with pytest.raises(AssertionError, match="duplicate scene keys"): @@ -161,13 +161,13 @@ def test_static_embodiment_placement_stores_per_env_poses(): from isaaclab_arena.environments.relation_solver_interface import _apply_relation_placement_result from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import PosePerEnv desk = _make_desk() robot = DummyEmbodiment( name="robot", - bounding_box=AxisAlignedBoundingBox( + bounding_box=OrientedBoundingBox.from_min_max( min_point=(-0.2, -0.2, 0.0), max_point=(0.2, 0.2, 1.0), ), @@ -228,7 +228,7 @@ def test_relation_placement_rejects_movable_asset_with_unplaced_auxiliary_prims( desk = _make_desk() box = _CompoundObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=0.01)) @@ -248,7 +248,7 @@ def test_anchor_asset_with_unplaced_auxiliary_prims_is_allowed(): anchor = _CompoundObject( name="anchored_stand", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), ) anchor.add_relation(IsAnchor()) # Anchors are excluded from the guard; this must not raise. diff --git a/isaaclab_arena/tests/test_usd_helpers.py b/isaaclab_arena/tests/test_usd_helpers.py index 522ef0b73e..30cbb99bd5 100644 --- a/isaaclab_arena/tests/test_usd_helpers.py +++ b/isaaclab_arena/tests/test_usd_helpers.py @@ -33,7 +33,7 @@ def _bbox_size(path: pathlib.Path, scale: tuple[float, float, float]) -> tuple[f from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd bbox = compute_local_bounding_box_from_usd(path.as_posix(), scale=scale) - size = bbox.size[0] + size = 2.0 * bbox.half_extents[0] return (float(size[0]), float(size[1]), float(size[2])) @@ -47,7 +47,7 @@ def _test_compute_local_bounding_box_from_usd(simulation_app, asset_dir: pathlib size = _bbox_size(unit_cube, scale=(2.0, 2.0, 2.0)) assert all(abs(dim - 2.0) < EPS for dim in size), size - # Default-prim root scale is unbaked before spawn scale is applied. + # ComputeUntransformedBound excludes the authored default-prim transform before spawn scale is applied. size = _bbox_size(scaled_root_cube, scale=(1.0, 1.0, 1.0)) assert all(abs(dim - 1.0) < EPS for dim in size), size diff --git a/isaaclab_arena/tests/test_usd_pose_helpers.py b/isaaclab_arena/tests/test_usd_pose_helpers.py index 442bead031..9c2944aac9 100644 --- a/isaaclab_arena/tests/test_usd_pose_helpers.py +++ b/isaaclab_arena/tests/test_usd_pose_helpers.py @@ -5,6 +5,8 @@ import numpy as np +import pytest + from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function HEADLESS = True @@ -46,6 +48,23 @@ def _test_get_prim_pose_in_default_prim_frame(simulation_app): return True +def _test_get_prim_pose_rejects_scaled_reference(simulation_app): + """A referenced prim with local scale cannot be represented by Pose and OBB.""" + from pxr import Gf, Usd, UsdGeom + + from isaaclab_arena.utils.usd_pose_helpers import get_prim_pose_in_default_prim_frame + + stage = Usd.Stage.CreateInMemory() + root = UsdGeom.Xform.Define(stage, "/Root") + stage.SetDefaultPrim(root.GetPrim()) + reference = UsdGeom.Xform.Define(stage, "/Root/Reference") + reference.AddScaleOp().Set(Gf.Vec3d(2.0, 2.0, 2.0)) + + with pytest.raises(AssertionError, match="must have unit scale"): + get_prim_pose_in_default_prim_frame(reference.GetPrim(), stage) + return True + + def test_get_prim_pose_in_default_prim_frame(): # Basic test that just adds all our pick-up objects to the scene and checks that nothing crashes. result = run_simulation_app_function( @@ -55,5 +74,10 @@ def test_get_prim_pose_in_default_prim_frame(): assert result, "Test failed" +def test_get_prim_pose_rejects_scaled_reference(): + """Scaled referenced prims fail before their local scale can be discarded.""" + assert run_simulation_app_function(_test_get_prim_pose_rejects_scaled_reference, headless=HEADLESS) + + if __name__ == "__main__": test_get_prim_pose_in_default_prim_frame() diff --git a/isaaclab_arena/tests/test_usd_scale_helpers.py b/isaaclab_arena/tests/test_usd_scale_helpers.py index ee3fb44351..f04ae7857c 100644 --- a/isaaclab_arena/tests/test_usd_scale_helpers.py +++ b/isaaclab_arena/tests/test_usd_scale_helpers.py @@ -5,8 +5,8 @@ """Regression tests for USD scale handling in extract_trimesh_from_usd and compute_local_bounding_box_from_usd. -Verifies that spawn-scale is applied in the local frame (R·(S·v)+t) rather than world -frame (S·(R·v+t)), which matters for translated/rotated child prims under non-uniform scale. +Verifies that spawn scale is applied in the default-prim frame after child transforms, +matching scale authored on the spawned default-prim wrapper. """ import numpy as np @@ -20,17 +20,16 @@ def _test_extract_trimesh_translated_child_nonuniform_scale(simulation_app): - """extract_trimesh_from_usd must scale in local frame, not world frame. + """Mesh and bounding-box extraction agree in the scaled default-prim frame. Setup: unit cube under a child Xform translated +1.0 in X, scale=(2,1,1). - Correct (local scale): verts ±0.5 → ±1.0 in local X, then translate +1 → world X [0.0, 2.0]. - Bug (world scale): verts ±0.5, translate +1 → world [0.5, 1.5], then *2 → [1.0, 3.0]. + Child transform gives X [0.5, 1.5], then root-wrapper scale gives X [1.0, 3.0]. """ import tempfile from pxr import Gf, Usd, UsdGeom - from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd + from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd, extract_trimesh_from_usd stage = Usd.Stage.CreateInMemory() root = stage.DefinePrim("/root", "Xform") @@ -87,17 +86,21 @@ def _test_extract_trimesh_translated_child_nonuniform_scale(simulation_app): scale = (2.0, 1.0, 1.0) tri = extract_trimesh_from_usd(usd_path, scale=scale) + bbox = compute_local_bounding_box_from_usd(usd_path, scale=scale) verts = tri.vertices - # Local scale: ±0.5*2=±1.0 then +1 translate → world X [0.0, 2.0] - assert np.isclose(verts[:, 0].min(), 0.0, atol=1e-5), f"got {verts[:, 0].min():.4f}" - assert np.isclose(verts[:, 0].max(), 2.0, atol=1e-5), f"got {verts[:, 0].max():.4f}" + assert np.isclose(verts[:, 0].min(), 1.0, atol=1e-5), f"got {verts[:, 0].min():.4f}" + assert np.isclose(verts[:, 0].max(), 3.0, atol=1e-5), f"got {verts[:, 0].max():.4f}" assert np.isclose(verts[:, 1].min(), -0.5, atol=1e-5) assert np.isclose(verts[:, 1].max(), 0.5, atol=1e-5) assert np.isclose(verts[:, 2].min(), -0.5, atol=1e-5) assert np.isclose(verts[:, 2].max(), 0.5, atol=1e-5) + bbox_min, bbox_max = bbox.get_axis_aligned_bounds() + np.testing.assert_allclose(bbox_min[0].cpu().numpy(), verts.min(axis=0), atol=1e-5) + np.testing.assert_allclose(bbox_max[0].cpu().numpy(), verts.max(axis=0), atol=1e-5) + return True @@ -332,8 +335,8 @@ def _test_bbox_translated_child_nonuniform_scale(simulation_app): bbox = compute_local_bounding_box_from_usd(usd_path, scale=scale) # ComputeLocalBound gives [0.5,1.5] * scale_x=2 → [1.0, 3.0] - min_pt = bbox.min_point[0] # (3,) tensor - max_pt = bbox.max_point[0] # (3,) tensor + min_point, max_point = bbox.get_axis_aligned_bounds() + min_pt, max_pt = min_point[0], max_point[0] assert np.isclose(min_pt[0].item(), 1.0, atol=1e-5), f"got {min_pt[0].item():.4f}" assert np.isclose(max_pt[0].item(), 3.0, atol=1e-5), f"got {max_pt[0].item():.4f}" assert np.isclose(min_pt[1].item(), -0.5, atol=1e-5) @@ -415,8 +418,8 @@ def _test_both_paths_agree_origin_prim(simulation_app): assert np.isclose(verts[:, 2].max(), 0.25, atol=1e-5) # BBox must match mesh extents exactly for origin-centered single prim. - min_pt = bbox.min_point[0] # (3,) tensor - max_pt = bbox.max_point[0] # (3,) tensor + min_point, max_point = bbox.get_axis_aligned_bounds() + min_pt, max_pt = min_point[0], max_point[0] assert np.isclose(min_pt[0].item(), -1.0, atol=1e-5) assert np.isclose(max_pt[0].item(), 1.0, atol=1e-5) assert np.isclose(min_pt[1].item(), -1.5, atol=1e-5) diff --git a/isaaclab_arena/tests/test_validate_placement.py b/isaaclab_arena/tests/test_validate_placement.py index 0f8f59d340..5a45c3e619 100644 --- a/isaaclab_arena/tests/test_validate_placement.py +++ b/isaaclab_arena/tests/test_validate_placement.py @@ -14,28 +14,28 @@ from isaaclab_arena.relations.placement_validators import NextToValidator, NotNextToValidator, OnRelationValidator from isaaclab_arena.relations.relations import NextTo, NotNextTo, On, RotateAroundSolution, Side from isaaclab_arena.tests.dummy_object import DummyObject -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox def _make_box(name: str, size: float = 0.2) -> DummyObject: half = size / 2 return DummyObject( name=name, - bounding_box=AxisAlignedBoundingBox(min_point=(-half, -half, -half), max_point=(half, half, half)), + bounding_box=OrientedBoundingBox.from_min_max((-half, -half, -half), (half, half, half)), ) def _make_long_box(name: str, half_x: float = 0.3, half_y: float = 0.05, half_z: float = 0.05) -> DummyObject: return DummyObject( name=name, - bounding_box=AxisAlignedBoundingBox(min_point=(-half_x, -half_y, -half_z), max_point=(half_x, half_y, half_z)), + bounding_box=OrientedBoundingBox.from_min_max((-half_x, -half_y, -half_z), (half_x, half_y, half_z)), ) def _make_desk() -> DummyObject: return DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(-0.5, -0.5, 0.0), max_point=(0.5, 0.5, 0.05)), + bounding_box=OrientedBoundingBox.from_min_max((-0.5, -0.5, 0.0), (0.5, 0.5, 0.05)), ) @@ -43,14 +43,18 @@ def _env_bboxes(positions: dict[DummyObject, tuple[float, float, float]]): return {obj: obj.get_bounding_box() for obj in positions} -def _validate_one(placer: ObjectPlacer, positions, env_bboxes, orientations=None): +def _validate_one(placer: ObjectPlacer, positions, env_bboxes, rotations=None): """Run every enabled validator over a single candidate and return its aggregated results.""" - return placer._validate_candidates([positions], [orientations or {}], [env_bboxes], [])[0] + return placer._validate_candidates([positions], [rotations or {}], [env_bboxes], [])[0] -def _stack_rows(bbox: AxisAlignedBoundingBox, n: int) -> AxisAlignedBoundingBox: +def _stack_rows(bbox: OrientedBoundingBox, n: int) -> OrientedBoundingBox: """Repeat a single-env bbox into n stacked rows (one per candidate).""" - return AxisAlignedBoundingBox(min_point=bbox.min_point.repeat(n, 1), max_point=bbox.max_point.repeat(n, 1)) + return OrientedBoundingBox( + bbox.center.repeat(n, 1), + bbox.half_extents.repeat(n, 1), + bbox.rotation_xyzw.repeat(n, 1), + ) def test_no_overlap_returns_true(): @@ -117,8 +121,8 @@ def test_rotation_aware_overlap_uses_yaw(): positions = {a: (0.0, 0.0, 0.0), b: (0.0, 0.2, 0.0)} axis_aligned = {a: a.get_bounding_box(), b: b.get_bounding_box()} assert _validate_one(placer, positions, axis_aligned).do_all_required_validation_checks_pass() is True - rotated = {a: a.get_bounding_box().rotated_around_z(math.pi / 2), b: b.get_bounding_box()} - assert _validate_one(placer, positions, rotated).do_all_required_validation_checks_pass() is False + rotations = {a: (0.0, 0.0, 2**-0.5, 2**-0.5)} + assert _validate_one(placer, positions, axis_aligned, rotations).do_all_required_validation_checks_pass() is False def test_candidate_bbox_aligns_with_candidate_yaw(): @@ -130,12 +134,14 @@ def test_candidate_bbox_aligns_with_candidate_yaw(): # Two candidates share positions but assign distinct yaws to `a`. candidate_bboxes = {a: _stack_rows(a.get_bounding_box(), 2), b: _stack_rows(b.get_bounding_box(), 2)} - rotated = ObjectPlacer._rotate_candidate_bboxes([a, b], candidate_bboxes, [{a: 0.0}, {a: math.pi / 2}]) - # Mirrors _place_ranked: each candidate validates against its own bbox row. + rotations = [{a: (0.0, 0.0, 0.0, 1.0)}, {a: (0.0, 0.0, 2**-0.5, 2**-0.5)}] validations = [ _validate_one( - placer, positions, ObjectPlacer._get_bounding_boxes_for_candidate_index(rotated, idx) + placer, + positions, + ObjectPlacer._get_bounding_boxes_for_candidate_index(candidate_bboxes, idx), + rotations[idx], ).do_all_required_validation_checks_pass() for idx in range(2) ] @@ -143,30 +149,29 @@ def test_candidate_bbox_aligns_with_candidate_yaw(): assert validations == [True, False] -def test_rotate_candidate_bboxes_encloses_marker_plus_sampled_yaw(): - """Rotated bbox equals the original bbox rotated by the combined marker+sampled yaw.""" - box = _make_long_box("box") - marker_yaw, sampled_yaw = math.pi / 6, math.pi / 3 - total_yaw = marker_yaw + sampled_yaw - box.add_relation(RotateAroundSolution(yaw_rad=marker_yaw)) +def test_candidate_rotation_composes_world_yaw_over_pitched_marker(monkeypatch): + """A forced sampled world yaw rotates independently expected pitched corners.""" + box = _make_long_box("box", half_x=0.3, half_y=0.1, half_z=0.05) + box.add_relation(RotateAroundSolution(pitch_rad=math.pi / 2)) + monkeypatch.setattr("isaaclab_arena.relations.object_placer.get_random_rotation", lambda generator: math.pi / 2) - rotated = ObjectPlacer._rotate_candidate_bboxes([box], {box: box.get_bounding_box()}, [{box: total_yaw}]) + rotation = ObjectPlacer(ObjectPlacerParams(random_yaw_init=True))._generate_initial_rotations([box], set())[box] + actual = box.get_bounding_box().rotated_by_quat(rotation).get_corners()[0] - expected = box.get_bounding_box().rotated_around_z(total_yaw) - torch.testing.assert_close(rotated[box].min_point, expected.min_point, atol=1e-6, rtol=0) - torch.testing.assert_close(rotated[box].max_point, expected.max_point, atol=1e-6, rtol=0) - # Passing only sampled_yaw (without marker) would enclose an undersized, misaligned footprint. - sampled_only = box.get_bounding_box().rotated_around_z(sampled_yaw) - assert not torch.allclose(rotated[box].max_point, sampled_only.max_point, atol=1e-6) + local = box.get_bounding_box().get_corners()[0] + # Rz(90) @ Ry(90): (x, y, z) -> (-y, z, -x). + expected = torch.stack((-local[:, 1], local[:, 2], -local[:, 0]), dim=1) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=0) def test_enclosing_after_rotation_pitch_swaps_extents(): """A 90° pitch rotates the tall Z extent into X, so the enclosing AABB swaps X and Z half-sizes.""" - box = AxisAlignedBoundingBox(min_point=(-0.05, -0.05, -0.3), max_point=(0.05, 0.05, 0.3)) + box = OrientedBoundingBox.from_min_max((-0.05, -0.05, -0.3), (0.05, 0.05, 0.3)) quat = RotateAroundSolution(pitch_rad=math.pi / 2).get_rotation_xyzw() - rotated = box.enclosing_after_rotation(quat) - torch.testing.assert_close(rotated.max_point, torch.tensor([[0.3, 0.05, 0.05]]), atol=1e-6, rtol=0) - torch.testing.assert_close(rotated.min_point, torch.tensor([[-0.3, -0.05, -0.05]]), atol=1e-6, rtol=0) + rotated = box.rotated_by_quat(quat) + minimum, maximum = rotated.get_axis_aligned_bounds() + torch.testing.assert_close(maximum, torch.tensor([[0.3, 0.05, 0.05]]), atol=1e-6, rtol=0) + torch.testing.assert_close(minimum, torch.tensor([[-0.3, -0.05, -0.05]]), atol=1e-6, rtol=0) def test_rotate_candidate_bboxes_encloses_pitched_object(): @@ -182,8 +187,8 @@ def test_rotate_candidate_bboxes_encloses_pitched_object(): assert _validate_one(placer, positions, axis_aligned).do_all_required_validation_checks_pass() is True # Composing the applied pitch grows a's X extent to 0.3, so it now overlaps b and is rejected. - rotated = ObjectPlacer._rotate_candidate_bboxes([a, b], axis_aligned, [{}]) - assert _validate_one(placer, positions, rotated).do_all_required_validation_checks_pass() is False + rotations = {a: RotateAroundSolution(pitch_rad=math.pi / 2).get_rotation_xyzw()} + assert _validate_one(placer, positions, axis_aligned, rotations).do_all_required_validation_checks_pass() is False def test_on_relation_containment_uses_rotated_bbox(): @@ -197,8 +202,8 @@ def test_on_relation_containment_uses_rotated_bbox(): axis_aligned = {desk: desk.get_bounding_box(), child: child.get_bounding_box()} assert OnRelationValidator(placer.params)._validate(positions, axis_aligned) is True - rotated = {desk: desk.get_bounding_box(), child: child.get_bounding_box().rotated_around_z(math.pi / 2)} - assert OnRelationValidator(placer.params)._validate(positions, rotated) is False + rotations = {child: (0.0, 0.0, 2**-0.5, 2**-0.5)} + assert OnRelationValidator(placer.params)._validate(positions, axis_aligned, rotations) is False def test_on_relation_check_no_relation_returns_true(): diff --git a/isaaclab_arena/utils/bounding_box.py b/isaaclab_arena/utils/bounding_box.py index bac96332fa..380d102c02 100644 --- a/isaaclab_arena/utils/bounding_box.py +++ b/isaaclab_arena/utils/bounding_box.py @@ -3,396 +3,305 @@ # # SPDX-License-Identifier: Apache-2.0 -""" -Utilities for computing spatial relationships between objects in a scene. -This module provides functions for: -- Computing bounding boxes from USD assets -- Calculating placement poses based on semantic relationships (e.g., "on_top_of") -- Supporting randomized placement with specified constraints -- AxisAlignedBoundingBox with always-tensor storage supporting single (N=1) and batched (N>1) modes -""" +"""Oriented bounding-box geometry utilities.""" import torch +from dataclasses import dataclass + +from isaaclab.utils.math import matrix_from_quat, quat_apply, quat_mul from isaaclab_arena.utils.pose import Pose +Vector3 = tuple[float, float, float] | torch.Tensor +Quaternion = tuple[float, float, float, float] | torch.Tensor -class AxisAlignedBoundingBox: - """Axis-aligned bounding box storing local extents. Use get_corners_at(pos) for world-space corners. - Stores min/max extents as (N, 3) float32 tensors where N is the number of environments. - All properties consistently return tensors: (N, 3) for point values, (N,) for scalars. - Constructor accepts tuples, 1D tensors, or (N, 3) tensors. - """ +def _as_batched(value: tuple[float, ...] | torch.Tensor, width: int, device: torch.device) -> torch.Tensor: + tensor = torch.as_tensor(value, dtype=torch.float32, device=device) + assert tensor.ndim in (1, 2), f"Expected a ({width},) or (N, {width}) value, got {tuple(tensor.shape)}." + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + assert tensor.shape[1] == width, f"Expected trailing dimension {width}, got {tuple(tensor.shape)}." + assert tensor.shape[0] > 0, "Bounding-box batches must not be empty." + return tensor - def __init__( - self, - min_point: tuple[float, float, float] | torch.Tensor, - max_point: tuple[float, float, float] | torch.Tensor, - ): - self._min_point = self._to_batched_tensor(min_point) - self._max_point = self._to_batched_tensor(max_point) - assert self._min_point.shape == self._max_point.shape - assert self._min_point.shape[-1] == 3 - - def __repr__(self) -> str: - return f"AxisAlignedBoundingBox(min_point={self._min_point}, max_point={self._max_point})" - - def __getitem__(self, idx: int) -> "AxisAlignedBoundingBox": - """Select row idx (one env/candidate), returning a single-box (N=1) bbox.""" - assert 0 <= idx < self.num_envs, f"index {idx} out of range for bbox with num_envs={self.num_envs}." - return AxisAlignedBoundingBox( - min_point=self._min_point[idx : idx + 1], max_point=self._max_point[idx : idx + 1] - ) - @staticmethod - def _to_batched_tensor(value: tuple[float, float, float] | torch.Tensor) -> torch.Tensor: - """Convert tuple, 1-D tensor, or (N, 3) tensor to (N, 3) float32 tensor.""" - if isinstance(value, tuple): - return torch.tensor([value], dtype=torch.float32) - if value.dim() == 1: - return value.unsqueeze(0).float() - return value.float() +def _common_device(*values: tuple[float, ...] | torch.Tensor) -> torch.device: + devices = [value.device for value in values if not isinstance(value, tuple)] + if not devices: + return torch.device("cpu") + assert all(device == devices[0] for device in devices), "Bounding-box inputs must be on the same device." + return devices[0] - @property - def min_point(self) -> torch.Tensor: - """Local minimum extent (x, y, z) relative to object origin. Shape (N, 3).""" - return self._min_point - @property - def max_point(self) -> torch.Tensor: - """Local maximum extent (x, y, z) relative to object origin. Shape (N, 3).""" - return self._max_point +def _broadcast_rows(*values: torch.Tensor) -> tuple[torch.Tensor, ...]: + count = max(value.shape[0] for value in values) + assert all( + value.shape[0] in (1, count) for value in values + ), f"Batched values must have equal leading dimensions or N=1; got {[value.shape[0] for value in values]}." + return tuple(value.expand(count, *value.shape[1:]) if value.shape[0] == 1 else value for value in values) - @property - def num_envs(self) -> int: - """Number of environments (leading dimension N).""" - return self._min_point.shape[0] - def is_batch_invariant(self) -> bool: - """Return True when every env shares the same bbox extents.""" - return bool( - torch.allclose(self._min_point, self._min_point[:1].expand_as(self._min_point)) - and torch.allclose(self._max_point, self._max_point[:1].expand_as(self._max_point)) - ) +@dataclass(init=False, slots=True, eq=False) +class OrientedBoundingBox: + """A batch of oriented boxes, where N is the number of boxes.""" - @property - def size(self) -> torch.Tensor: - """Returns the size (width, depth, height) of the bounding box. Shape (N, 3).""" - return self._max_point - self._min_point + center: torch.Tensor + """Box centers. Shape (N, 3).""" - @property - def center(self) -> torch.Tensor: - """Returns the center point of the bounding box. Shape (N, 3).""" - return (self._min_point + self._max_point) * 0.5 + half_extents: torch.Tensor + """Non-negative half-lengths along each local box axis. Shape (N, 3).""" - @property - def top_surface_z(self) -> torch.Tensor: - """Returns the z-coordinate of the top surface. Shape (N,).""" - return self._max_point[:, 2] - - @property - def bottom_surface_z(self) -> torch.Tensor: - """Returns the z-coordinate of the bottom surface. Shape (N,).""" - return self._min_point[:, 2] - - def get_corners_at(self, pos: torch.Tensor | None = None) -> torch.Tensor: - """Get 8 corners of this bounding box, optionally offset by position. - - Args: - pos: If provided, world position (x, y, z) to offset corners by. - If None, returns corners in local/object frame. + rotation_xyzw: torch.Tensor + """Local-to-parent unit quaternions in xyzw order. Shape (N, 4).""" - Returns: - Tensor of shape (N, 8, 3) with corners ordered: bottom 4, then top 4. - """ - min_pt, max_pt = self._min_point, self._max_point - corners = torch.stack( - [ - torch.stack([min_pt[:, 0], min_pt[:, 1], min_pt[:, 2]], dim=1), # Bottom-front-left - torch.stack([max_pt[:, 0], min_pt[:, 1], min_pt[:, 2]], dim=1), # Bottom-front-right - torch.stack([max_pt[:, 0], max_pt[:, 1], min_pt[:, 2]], dim=1), # Bottom-back-right - torch.stack([min_pt[:, 0], max_pt[:, 1], min_pt[:, 2]], dim=1), # Bottom-back-left - torch.stack([min_pt[:, 0], min_pt[:, 1], max_pt[:, 2]], dim=1), # Top-front-left - torch.stack([max_pt[:, 0], min_pt[:, 1], max_pt[:, 2]], dim=1), # Top-front-right - torch.stack([max_pt[:, 0], max_pt[:, 1], max_pt[:, 2]], dim=1), # Top-back-right - torch.stack([min_pt[:, 0], max_pt[:, 1], max_pt[:, 2]], dim=1), # Top-back-left - ], - dim=1, + def __init__(self, center: Vector3, half_extents: Vector3, rotation_xyzw: Quaternion): + device = _common_device(center, half_extents, rotation_xyzw) + center_tensor = _as_batched(center, 3, device) + half_extents_tensor = _as_batched(half_extents, 3, device) + rotation_tensor = _as_batched(rotation_xyzw, 4, device) + center_tensor, half_extents_tensor, rotation_tensor = _broadcast_rows( + center_tensor, half_extents_tensor, rotation_tensor ) - if pos is not None: - if pos.dim() == 1: - pos = pos.unsqueeze(0) - corners = corners + pos.unsqueeze(1) - return corners - - def scaled(self, scale: tuple[float, float, float] | torch.Tensor) -> "AxisAlignedBoundingBox": - """Return a new bounding box with scale applied. - - Args: - scale: Scale factors (x, y, z) to apply. - - Returns: - New AxisAlignedBoundingBox with scaled dimensions. - """ - scale = self._to_batched_tensor(scale) - return AxisAlignedBoundingBox(min_point=self._min_point * scale, max_point=self._max_point * scale) - - def to(self, device: torch.device) -> "AxisAlignedBoundingBox": - """Return a new bounding box with tensors on *device*.""" - return AxisAlignedBoundingBox(min_point=self._min_point.to(device), max_point=self._max_point.to(device)) - - def translated(self, offset: tuple[float, float, float] | torch.Tensor) -> "AxisAlignedBoundingBox": - """Return a new bounding box translated by an offset. - - Args: - offset: Translation offset (x, y, z) to apply. - - Returns: - New AxisAlignedBoundingBox with translated position. - """ - offset = self._to_batched_tensor(offset) - return AxisAlignedBoundingBox(min_point=self._min_point + offset, max_point=self._max_point + offset) - - def centered(self) -> "AxisAlignedBoundingBox": - """Return a new bounding box centered around the origin. - - The returned bbox has the same size but is shifted so that its - center is at (0, 0, 0). - - Returns: - New AxisAlignedBoundingBox centered at origin. - """ - center = (self._min_point + self._max_point) * 0.5 - return AxisAlignedBoundingBox(min_point=self._min_point - center, max_point=self._max_point - center) - - def overlaps(self, other: "AxisAlignedBoundingBox", margin: float = 0.0) -> torch.Tensor: - """Check if two AABBs overlap in 3D. - - Args: - other: The other bounding box to test against. - margin: Minimum required separation in meters. A positive value - rejects placements where the gap is smaller than margin. - - Returns: - Bool tensor of shape (N,). True where volumes overlap (or are closer than margin). - """ - return ( - (self._max_point[:, 0] + margin > other._min_point[:, 0]) - & (other._max_point[:, 0] + margin > self._min_point[:, 0]) - & (self._max_point[:, 1] + margin > other._min_point[:, 1]) - & (other._max_point[:, 1] + margin > self._min_point[:, 1]) - & (self._max_point[:, 2] + margin > other._min_point[:, 2]) - & (other._max_point[:, 2] + margin > self._min_point[:, 2]) + center_tensor = center_tensor.clone() + half_extents_tensor = half_extents_tensor.clone() + rotation_tensor = rotation_tensor.clone() + + assert torch.isfinite(center_tensor).all(), "Box centers must be finite." + assert torch.isfinite(half_extents_tensor).all(), "Box half-extents must be finite." + assert (half_extents_tensor >= 0.0).all(), "Box half-extents must be non-negative." + assert torch.isfinite(rotation_tensor).all(), "Box rotations must be finite." + norms = torch.linalg.vector_norm(rotation_tensor, dim=-1) + assert torch.allclose( + norms, torch.ones_like(norms), atol=1e-5, rtol=1e-5 + ), "Box rotations must be unit quaternions." + + self.center = center_tensor + self.half_extents = half_extents_tensor + self.rotation_xyzw = rotation_tensor + + @classmethod + def from_tensors_unchecked( + cls, + center: torch.Tensor, + half_extents: torch.Tensor, + rotation_xyzw: torch.Tensor, + ) -> "OrientedBoundingBox": + """Build from validated tensors without cloning or synchronizing the device.""" + bbox = cls.__new__(cls) + bbox.center = center + bbox.half_extents = half_extents + bbox.rotation_xyzw = rotation_xyzw + return bbox + + @classmethod + def from_min_max(cls, min_point: Vector3, max_point: Vector3) -> "OrientedBoundingBox": + """Construct axis-aligned boxes from minimum and maximum points.""" + device = _common_device(min_point, max_point) + minimum = _as_batched(min_point, 3, device) + maximum = _as_batched(max_point, 3, device) + minimum, maximum = _broadcast_rows(minimum, maximum) + assert (maximum >= minimum).all(), "Maximum points must not be below minimum points." + identity = minimum.new_zeros((1, 4)) + identity[:, 3] = 1.0 + identity = identity.expand(minimum.shape[0], 4) + return cls.from_tensors_unchecked((minimum + maximum) * 0.5, (maximum - minimum) * 0.5, identity) + + def __getitem__(self, idx: int) -> "OrientedBoundingBox": + """Select one box while preserving its leading dimension.""" + assert 0 <= idx < self.num_envs, f"Index {idx} out of range for {self.num_envs} boxes." + return self.from_tensors_unchecked( + self.center[idx : idx + 1], + self.half_extents[idx : idx + 1], + self.rotation_xyzw[idx : idx + 1], ) - def rotated_90_around_z(self, quarters: int) -> "AxisAlignedBoundingBox": - """Rotate AABB by quarters * 90° around Z axis. - - Only 90° increments are supported to preserve axis-alignment without size increase. - - Args: - quarters: Number of 90° rotations (0=0°, 1=90°, 2=180°, 3=270°/-90°). - - Returns: - New AxisAlignedBoundingBox rotated around Z axis. - """ - quarters = quarters % 4 - min_x, min_y, min_z = self._min_point[:, 0], self._min_point[:, 1], self._min_point[:, 2] - max_x, max_y, max_z = self._max_point[:, 0], self._max_point[:, 1], self._max_point[:, 2] - if quarters == 0: - return AxisAlignedBoundingBox(min_point=self._min_point.clone(), max_point=self._max_point.clone()) - elif quarters == 1: # 90° CCW - return AxisAlignedBoundingBox( - min_point=torch.stack([-max_y, min_x, min_z], dim=1), - max_point=torch.stack([-min_y, max_x, max_z], dim=1), - ) - elif quarters == 2: # 180° - return AxisAlignedBoundingBox( - min_point=torch.stack([-max_x, -max_y, min_z], dim=1), - max_point=torch.stack([-min_x, -min_y, max_z], dim=1), - ) - else: # 270° CCW / -90° (quarters == 3) - return AxisAlignedBoundingBox( - min_point=torch.stack([min_y, -max_x, min_z], dim=1), - max_point=torch.stack([max_y, -min_x, max_z], dim=1), - ) - - def enclosing_after_rotation( - self, rotation_xyzw: tuple[float, float, float, float] | torch.Tensor - ) -> "AxisAlignedBoundingBox": - """Refit to the axis-aligned box enclosing this box under an arbitrary rotation. - - Rotates the eight corners about the object origin by ``rotation_xyzw`` and returns the - tightest AABB containing them. Conservative (larger than the true rotated box) for any - non-axis-aligned rotation. Unlike :meth:`rotated_around_z`, roll and pitch tilt the box - out of plane, so the Z extent may grow as well. - - Args: - rotation_xyzw: Quaternion ``(x, y, z, w)`` applied about the object origin. A single - quaternion rotates every stacked box equally. + @property + def num_envs(self) -> int: + """Return the leading dimension N.""" + return self.center.shape[0] - Returns: - New AxisAlignedBoundingBox enclosing the rotated corners. - """ - device = self._min_point.device - quat = torch.as_tensor(rotation_xyzw, dtype=torch.float32, device=device).reshape(-1) - assert quat.shape == ( - 4, - ), f"enclosing_after_rotation expects a single (x, y, z, w) quaternion, got {tuple(quat.shape)}." - qx, qy, qz, qw = quat.unbind(0) - # Rotation matrix from the (x, y, z, w) quaternion. - rot = torch.stack([ - torch.stack([1 - 2 * (qy * qy + qz * qz), 2 * (qx * qy - qz * qw), 2 * (qx * qz + qy * qw)]), - torch.stack([2 * (qx * qy + qz * qw), 1 - 2 * (qx * qx + qz * qz), 2 * (qy * qz - qx * qw)]), - torch.stack([2 * (qx * qz - qy * qw), 2 * (qy * qz + qx * qw), 1 - 2 * (qx * qx + qy * qy)]), - ]) # (3, 3) - corners = self.get_corners_at() # (N, 8, 3) - rotated = corners @ rot.transpose(0, 1) # (N, 8, 3): each corner mapped by rot - return AxisAlignedBoundingBox( - min_point=rotated.min(dim=1).values, - max_point=rotated.max(dim=1).values, + def is_batch_invariant(self) -> bool: + """Return whether every row describes the same box.""" + return all( + torch.allclose(value, value[:1].expand_as(value)) + for value in (self.center, self.half_extents, self.rotation_xyzw) ) - def rotated_around_z(self, angle_rad: float | torch.Tensor) -> "AxisAlignedBoundingBox": - """Refit to the axis-aligned box enclosing this box rotated by angle_rad around Z. - - Conservative (larger than the true rotated box) except at 90° multiples; Z extents unchanged. - - Args: - angle_rad: Yaw in radians. A scalar rotates every box equally; a 1-D tensor gives - per-box angles (or, for a single box, one box per angle). + def to(self, device: torch.device | str) -> "OrientedBoundingBox": + """Return the boxes on the requested device.""" + return self.from_tensors_unchecked( + self.center.to(device), + self.half_extents.to(device), + self.rotation_xyzw.to(device), + ) - Returns: - New AxisAlignedBoundingBox enclosing the rotated box. - """ - device = self._min_point.device - angles = torch.as_tensor(angle_rad, dtype=torch.float32, device=device).reshape(-1) # (M,) + def translated(self, offset: Vector3) -> "OrientedBoundingBox": + """Return boxes translated in their parent frame.""" + offset_tensor = _as_batched(offset, 3, self.center.device) + center, half_extents, rotation, offset_tensor = _broadcast_rows( + self.center, self.half_extents, self.rotation_xyzw, offset_tensor + ) + return self.from_tensors_unchecked(center + offset_tensor, half_extents, rotation) + + def rotated_by_quat(self, rotation_xyzw: Quaternion) -> "OrientedBoundingBox": + """Rotate boxes about the parent-frame origin.""" + rotation = _as_batched(rotation_xyzw, 4, self.center.device) + norms = torch.linalg.vector_norm(rotation, dim=-1) + assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5, rtol=1e-5), "Rotation must be unit length." + return self.rotated_by_quat_unchecked(rotation) + + def rotated_by_quat_unchecked(self, rotation: torch.Tensor) -> "OrientedBoundingBox": + """Rotate by validated batched quaternions without synchronizing the device.""" + rotation = _as_batched(rotation, 4, self.center.device) + center, half_extents, box_rotation, rotation = _broadcast_rows( + self.center, self.half_extents, self.rotation_xyzw, rotation + ) + return self.from_tensors_unchecked( + quat_apply(rotation, center), + half_extents, + quat_mul(rotation, box_rotation), + ) - num_boxes, num_angles = self._min_point.shape[0], angles.shape[0] - assert num_boxes == 1 or num_angles == 1 or num_boxes == num_angles, ( - "rotated_around_z requires one box, one angle, or equal counts; " - f"got {num_boxes} boxes and {num_angles} angles." + def transformed(self, position_xyz: Vector3, rotation_xyzw: Quaternion) -> "OrientedBoundingBox": + """Apply a parent-frame rigid transform to the boxes.""" + position = _as_batched(position_xyz, 3, self.center.device) + rotation = _as_batched(rotation_xyzw, 4, self.center.device) + norms = torch.linalg.vector_norm(rotation, dim=-1) + assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5, rtol=1e-5), "Rotation must be unit length." + return self.transformed_unchecked(position, rotation) + + def transformed_unchecked(self, position: torch.Tensor, rotation: torch.Tensor) -> "OrientedBoundingBox": + """Apply validated batched transforms without synchronizing the device.""" + position = _as_batched(position, 3, self.center.device) + rotated = self.rotated_by_quat_unchecked(rotation) + center, half_extents, box_rotation, position = _broadcast_rows( + rotated.center, rotated.half_extents, rotated.rotation_xyzw, position + ) + return self.from_tensors_unchecked(center + position, half_extents, box_rotation) + + def get_corners(self) -> torch.Tensor: + """Return parent-frame corners with shape (N, 8, 3).""" + signs = self.center.new_tensor([ + [-1.0, -1.0, -1.0], + [1.0, -1.0, -1.0], + [1.0, 1.0, -1.0], + [-1.0, 1.0, -1.0], + [-1.0, -1.0, 1.0], + [1.0, -1.0, 1.0], + [1.0, 1.0, 1.0], + [-1.0, 1.0, 1.0], + ]) + local_corners = signs.unsqueeze(0) * self.half_extents.unsqueeze(1) + rotation = self.rotation_xyzw.unsqueeze(1).expand(-1, 8, -1).reshape(-1, 4) + rotated = quat_apply(rotation, local_corners.reshape(-1, 3)).reshape(-1, 8, 3) + return self.center.unsqueeze(1) + rotated + + def get_bounds_along_axis(self, axis: Vector3) -> tuple[torch.Tensor, torch.Tensor]: + """Return minimum and maximum projections onto an axis.""" + axis_tensor = _as_batched(axis, 3, self.center.device) + center, half_extents, rotation, axis_tensor = _broadcast_rows( + self.center, self.half_extents, self.rotation_xyzw, axis_tensor + ) + axis_norm = torch.linalg.vector_norm(axis_tensor, dim=-1, keepdim=True) + assert (axis_norm > 0.0).all(), "Projection axes must be non-zero." + axis_tensor = axis_tensor / axis_norm + local_axes = matrix_from_quat(rotation).transpose(1, 2) + radius = (half_extents * torch.abs(torch.einsum("nid,nd->ni", local_axes, axis_tensor))).sum(dim=-1) + projected_center = (center * axis_tensor).sum(dim=-1) + return projected_center - radius, projected_center + radius + + def get_axis_aligned_bounds(self) -> tuple[torch.Tensor, torch.Tensor]: + """Return enclosing axis-aligned minimum and maximum points.""" + rotation = matrix_from_quat(self.rotation_xyzw) + radius = torch.bmm(rotation.abs(), self.half_extents.unsqueeze(-1)).squeeze(-1) + return self.center - radius, self.center + radius + + def is_axis_aligned(self, atol: float = 1e-5) -> torch.Tensor: + """Return whether each orientation is axis-aligned, with shape (N,).""" + absolute = matrix_from_quat(self.rotation_xyzw).abs() + near_zero_or_one = (absolute <= atol) | ((absolute - 1.0).abs() <= atol) + ones = torch.ones_like(absolute[..., 0]) + return ( + near_zero_or_one.all(dim=(-1, -2)) + & torch.isclose(absolute.sum(dim=-1), ones, atol=atol, rtol=0).all(dim=-1) + & torch.isclose(absolute.sum(dim=-2), ones, atol=atol, rtol=0).all(dim=-1) ) - cos = torch.cos(angles).unsqueeze(1) # (M, 1) - sin = torch.sin(angles).unsqueeze(1) # (M, 1) - - # XY footprint corners relative to the object origin: (N, 4). - min_x, min_y = self._min_point[:, 0], self._min_point[:, 1] - max_x, max_y = self._max_point[:, 0], self._max_point[:, 1] - corners_x = torch.stack([min_x, max_x, max_x, min_x], dim=1) # (N, 4) - corners_y = torch.stack([min_y, min_y, max_y, max_y], dim=1) # (N, 4) - - # Rotate corners (broadcasts (N, 4) with (M, 1)). - rot_x = corners_x * cos - corners_y * sin - rot_y = corners_x * sin + corners_y * cos - - new_min_x = rot_x.min(dim=1).values # (L,) - new_max_x = rot_x.max(dim=1).values - new_min_y = rot_y.min(dim=1).values - new_max_y = rot_y.max(dim=1).values - - # Z extents are invariant under Z rotation; broadcast to the output leading dim. - out_len = new_min_x.shape[0] - min_z, max_z = self._min_point[:, 2], self._max_point[:, 2] - if min_z.shape[0] == 1 and out_len > 1: - min_z = min_z.expand(out_len) - max_z = max_z.expand(out_len) - - return AxisAlignedBoundingBox( - min_point=torch.stack([new_min_x, new_min_y, min_z], dim=1), - max_point=torch.stack([new_max_x, new_max_y, max_z], dim=1), + def _paired(self, other: "OrientedBoundingBox") -> tuple[torch.Tensor, ...]: + assert self.center.device == other.center.device, "Bounding boxes must be on the same device." + return _broadcast_rows( + self.center, + self.half_extents, + self.rotation_xyzw, + other.center, + other.half_extents, + other.rotation_xyzw, ) - def rotated_by_quat( - self, rotation_xyzw: tuple[float, float, float, float] | torch.Tensor - ) -> "AxisAlignedBoundingBox": - """Refit to the axis-aligned box enclosing this box rotated about its origin by a quaternion. + def penetration( + self, + other: "OrientedBoundingBox", + clearance_m: float = 0.0, + tie_break_sign: float | torch.Tensor = 1.0, + ) -> torch.Tensor: + """Return positive SAT penetration or clearance violation with shape (N,). Args: - rotation_xyzw: Rotation quaternion as (x, y, z, w). A single quaternion rotates every box - equally; a (M, 4) tensor gives per-box quaternions (or, for a single box, one box per - quaternion), matching rotated_around_z's batching. - """ - corners = self.get_corners_at() # (N, 8, 3) - quats = torch.as_tensor(rotation_xyzw, dtype=corners.dtype, device=corners.device).reshape(-1, 4) # (M, 4) + other: Boxes to test, broadcastable over N. + clearance_m: Additional required separation in metres. + tie_break_sign: Directed escape sign for exactly concentric projections. - num_boxes, num_quats = corners.shape[0], quats.shape[0] - assert ( - num_boxes == 1 or num_quats == 1 or num_boxes == num_quats - ), f"rotated_by_quat requires one box, one quat, or equal counts; got {num_boxes} boxes and {num_quats} quats." - out_len = max(num_boxes, num_quats) - if num_boxes == 1 and out_len > 1: - corners = corners.expand(out_len, 8, 3) - - # Rotate the 8 object-frame corners by the quaternion (v + 2w(a×v) + 2a×(a×v)), then min/max. - axis = quats[:, :3].view(num_quats, 1, 3).expand(out_len, 8, 3) # (L, 8, 3) - qw = quats[:, 3].view(num_quats, 1, 1).expand(out_len, 1, 1) # (L, 1, 1) - axis_cross = torch.linalg.cross(axis, corners, dim=-1) - rotated = corners + 2.0 * qw * axis_cross + 2.0 * torch.linalg.cross(axis, axis_cross, dim=-1) - return AxisAlignedBoundingBox(min_point=rotated.amin(dim=1), max_point=rotated.amax(dim=1)) - - -def quaternion_to_90_deg_z_quarters(rotation_xyzw: tuple[float, float, float, float], tol_deg: float = 1.0) -> int: - """Convert a quaternion to 90° rotation quarters around Z axis. - - Only supports rotations that are multiples of 90° around the Z axis. - Raises AssertionError for any other rotation. - - Args: - rotation_xyzw: Quaternion as (x, y, z, w). - tol_deg: Tolerance in degrees for how close the angle must be to a 90° multiple. - - Returns: - Number of 90° quarters (0, 1, 2, or 3). - - Raises: - AssertionError: If the quaternion is not a pure Z rotation or not a 90° multiple. - """ - import math - - x, y, z, w = rotation_xyzw - - # Must be a pure Z rotation (x and y components must be ~0) - assert ( - abs(x) < 1e-3 and abs(y) < 1e-3 - ), f"Only rotations around Z axis are supported. Got quaternion (w={w:.4f}, x={x:.4f}, y={y:.4f}, z={z:.4f})." - - # Compute rotation angle around Z and normalize to [0°, 360°) - angle_deg = math.degrees(2 * math.atan2(z, w)) % 360 - quarters = round(angle_deg / 90) % 4 - remainder_deg = min(angle_deg % 90, 90 - angle_deg % 90) - - assert remainder_deg < tol_deg, ( - "Only 90° rotation multiples around Z are supported. " - f"Got {angle_deg:.1f}° (nearest 90° multiple: {quarters * 90}°)." - ) - - return quarters - - -def get_random_pose_within_bounding_box(bbox: AxisAlignedBoundingBox, seed: int | None = None) -> Pose: - """Generate a random pose (position and identity rotation) with position uniformly - sampled within a bounding box. - - Args: - bbox: Bounding box defining the valid region for sampling - seed: Optional random seed for reproducibility - - Returns: - Pose with random position within bbox and identity rotation - """ + Rows with separated enclosing AABBs are broadphase-culled to zero. + """ + assert clearance_m >= 0.0, "Clearance must be non-negative." + center_a, extent_a, quat_a, center_b, extent_b, quat_b = self._paired(other) + count = center_a.shape[0] + if isinstance(tie_break_sign, float): + assert tie_break_sign in (-1.0, 1.0), "Tie-break sign must be -1 or 1." + tie_sign = torch.as_tensor(tie_break_sign, dtype=center_a.dtype, device=center_a.device).reshape(-1, 1) + assert tie_sign.shape[0] in (1, count), f"Expected one tie-break sign or N={count}, got {tie_sign.shape[0]}." + rotation_a = matrix_from_quat(quat_a) + rotation_b = matrix_from_quat(quat_b) + aabb_radius_a = torch.bmm(rotation_a.abs(), extent_a.unsqueeze(-1)).squeeze(-1) + aabb_radius_b = torch.bmm(rotation_b.abs(), extent_b.unsqueeze(-1)).squeeze(-1) + min_a, max_a = center_a - aabb_radius_a, center_a + aabb_radius_a + min_b, max_b = center_b - aabb_radius_b, center_b + aabb_radius_b + broadphase = ((max_a + clearance_m > min_b) & (max_b + clearance_m > min_a)).all(dim=-1) + + axes_a = rotation_a.transpose(1, 2) + axes_b = rotation_b.transpose(1, 2) + cross_axes = torch.linalg.cross(axes_a.unsqueeze(2), axes_b.unsqueeze(1), dim=-1).reshape(count, 9, 3) + axes = torch.cat([axes_a, axes_b, cross_axes], dim=1) + axis_norm = torch.linalg.vector_norm(axes, dim=-1, keepdim=True) + valid = axis_norm.squeeze(-1) > 1e-6 + axes = axes / axis_norm.clamp_min(1e-6) + dominant_index = torch.abs(axes).argmax(dim=-1, keepdim=True) + dominant_component = torch.gather(axes, dim=-1, index=dominant_index) + canonical_sign = torch.where(dominant_component < 0.0, -1.0, 1.0) + axes = axes * canonical_sign + + radius_a = (torch.abs(torch.einsum("nkd,njd->nkj", axes, axes_a)) * extent_a.unsqueeze(1)).sum(dim=-1) + radius_b = (torch.abs(torch.einsum("nkd,njd->nkj", axes, axes_b)) * extent_b.unsqueeze(1)).sum(dim=-1) + projection = torch.einsum("nkd,nd->nk", axes, center_b - center_a) + # abs() has zero derivative at zero. Preserve its value while selecting a directed + # derivative so exactly concentric movable boxes receive an escape gradient. + signed_zero_projection = tie_sign * projection + distance = torch.where(projection == 0.0, signed_zero_projection, torch.abs(projection)) + depths = radius_a + radius_b + clearance_m - distance + depths = torch.where(valid, depths, torch.full_like(depths, torch.inf)) + penetration = depths.amin(dim=-1).clamp_min(0.0) + return torch.where(broadphase, penetration, torch.zeros_like(penetration)) + + def overlaps(self, other: "OrientedBoundingBox", clearance_m: float = 0.0) -> torch.Tensor: + """Return whether boxes overlap or violate the requested clearance.""" + return self.penetration(other, clearance_m) > 0.0 + + +def get_random_pose_within_bounding_box(bbox: OrientedBoundingBox, seed: int | None = None) -> Pose: + """Sample a position uniformly within the first box and use identity rotation.""" if seed is not None: torch.manual_seed(seed) - - # Get workspace bounds as (3,) tensors from the first (and typically only) environment - min_point = bbox.min_point[0] - max_point = bbox.max_point[0] - - # Sample random position uniformly within workspace bounds - random_position = min_point + (max_point - min_point) * torch.rand(3) - - pose = Pose(position_xyz=tuple(random_position.tolist()), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) - - return pose + local_position = (2.0 * torch.rand(3, device=bbox.center.device) - 1.0) * bbox.half_extents[0] + position = bbox.center[0] + quat_apply(bbox.rotation_xyzw[:1], local_position.unsqueeze(0))[0] + return Pose(position_xyz=tuple(position.cpu().tolist()), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) diff --git a/isaaclab_arena/utils/isaac_sim_debug_draw.py b/isaaclab_arena/utils/isaac_sim_debug_draw.py index 95d808a5c2..41874f1b4a 100644 --- a/isaaclab_arena/utils/isaac_sim_debug_draw.py +++ b/isaaclab_arena/utils/isaac_sim_debug_draw.py @@ -55,7 +55,19 @@ def draw_bbox( max_point: Maximum corner (x, y, z). thickness: Line thickness in pixels. """ - self._draw_bbox_wireframe(min_point, max_point, DEFAULT_COLOR, thickness) + x0, y0, z0 = min_point + x1, y1, z1 = max_point + corners = [ + (x0, y0, z0), + (x1, y0, z0), + (x1, y1, z0), + (x0, y1, z0), + (x0, y0, z1), + (x1, y0, z1), + (x1, y1, z1), + (x0, y1, z1), + ] + self._draw_corner_wireframe(corners, DEFAULT_COLOR, thickness) def draw_object_bboxes( self, @@ -64,59 +76,28 @@ def draw_object_bboxes( ) -> None: """Draw bounding boxes for one or more objects. - Uses each object's get_world_bounding_box() method which returns - the bounding box in world coordinates (local bbox + position offset). + Uses each object's exact oriented world bounding box. Args: objects: List of objects with get_world_bounding_box() methods. thickness: Line thickness in pixels. """ for obj in objects: - bbox_coords = self._extract_bbox_from_object(obj) - if bbox_coords is not None: - min_pt, max_pt = bbox_coords - self._draw_bbox_wireframe(min_pt, max_pt, DEFAULT_COLOR, thickness) - else: - print(f"Skipping {obj.name}: no bbox coordinates") + corners = obj.get_world_bounding_box().get_corners()[0].tolist() + self._draw_corner_wireframe(corners, DEFAULT_COLOR, thickness) def clear(self) -> None: """Clear all debug drawings.""" self._draw.clear_lines() self._draw.clear_points() - def _extract_bbox_from_object(self, obj) -> tuple[tuple, tuple] | None: - """Extract world-space bounding box coordinates from an object. - - Returns: - Tuple of (min_point, max_point) or None if extraction failed. - """ - world_bbox = obj.get_world_bounding_box() - return tuple(world_bbox.min_point[0].tolist()), tuple(world_bbox.max_point[0].tolist()) - - def _draw_bbox_wireframe( + def _draw_corner_wireframe( self, - min_point: tuple[float, float, float], - max_point: tuple[float, float, float], + corners: list[tuple[float, float, float]] | list[list[float]], color: tuple[float, float, float, float], thickness: float, ) -> None: - """Draw a wireframe bounding box using 12 edge lines.""" - x0, y0, z0 = min_point - x1, y1, z1 = max_point - - # 8 corners of the box - corners = [ - (x0, y0, z0), - (x1, y0, z0), - (x1, y1, z0), - (x0, y1, z0), # Bottom face - (x0, y0, z1), - (x1, y0, z1), - (x1, y1, z1), - (x0, y1, z1), # Top face - ] - - # 12 edges connecting corners + """Draw the 12-edge wireframe joining eight ordered corners.""" edges = [ (0, 1), (1, 2), @@ -132,7 +113,6 @@ def _draw_bbox_wireframe( (3, 7), # Vertical edges ] - # Build lists for draw_lines API start_points = [corners[i] for i, j in edges] end_points = [corners[j] for i, j in edges] colors_list = [color] * len(edges) diff --git a/isaaclab_arena/utils/trimesh.py b/isaaclab_arena/utils/trimesh.py index bd738bd90e..e6aa8bb70e 100644 --- a/isaaclab_arena/utils/trimesh.py +++ b/isaaclab_arena/utils/trimesh.py @@ -10,7 +10,7 @@ import numpy as np import trimesh -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -24,10 +24,10 @@ def mesh_in_world_frame(mesh: trimesh.Trimesh, pose: Pose) -> trimesh.Trimesh: return transformed -def bounding_box_from_mesh(mesh: trimesh.Trimesh) -> AxisAlignedBoundingBox: +def bounding_box_from_mesh(mesh: trimesh.Trimesh) -> OrientedBoundingBox: """Return the axis-aligned bounds of mesh.""" bounds = mesh.bounds - return AxisAlignedBoundingBox( + return OrientedBoundingBox.from_min_max( min_point=tuple(float(v) for v in bounds[0]), max_point=tuple(float(v) for v in bounds[1]), ) diff --git a/isaaclab_arena/utils/usd_helpers.py b/isaaclab_arena/utils/usd_helpers.py index b942a631b1..7d124823b5 100644 --- a/isaaclab_arena/utils/usd_helpers.py +++ b/isaaclab_arena/utils/usd_helpers.py @@ -9,10 +9,10 @@ import trimesh from contextlib import contextmanager -from pxr import Gf, Usd, UsdGeom, UsdLux, UsdPhysics +from pxr import Usd, UsdGeom, UsdLux, UsdPhysics from isaaclab_arena.assets.object_type import ObjectType -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox class NoCollisionMeshError(ValueError): @@ -157,37 +157,21 @@ def get_asset_usd_path_from_prim_path(prim_path: str, stage: Usd.Stage) -> str | return None -def _read_default_prim_scale(prim: Usd.Prim) -> tuple[float, float, float]: - """Return the default prim's root ``xformOp:scale``, or identity if absent.""" - if not prim.IsA(UsdGeom.Xformable): - return (1.0, 1.0, 1.0) - for op in UsdGeom.Xformable(prim).GetOrderedXformOps(): - if op.GetOpType() == UsdGeom.XformOp.TypeScale: - value = op.Get() - if value is not None: - return (float(value[0]), float(value[1]), float(value[2])) - return (1.0, 1.0, 1.0) - - def compute_local_bounding_box_from_usd( usd_path: str, scale: tuple[float, float, float] = (1.0, 1.0, 1.0), -) -> AxisAlignedBoundingBox: - """Compute the local bounding box matching Isaac Lab ``UsdFileCfg`` spawn size. +) -> OrientedBoundingBox: + """Compute default-prim-local bounds matching an Isaac Lab USD spawn. - Opening a USD directly includes the default prim's root ``xformOp:scale`` - in ``ComputeWorldBound``, but Isaac Lab's spawner ignores it and only - Object.scale on the spawn wrapper applies. - This helper unbakes the default prim's root scale from the USD, then - applies ``Object.scale`` once so relation-solver bboxes match what is - actually spawned. + The default prim's own transform is excluded because Isaac Lab's spawner + ignores it. ``scale`` is applied once to match the spawn wrapper. Args: usd_path: Path to the USD file. scale: Spawn-time scale passed to ``UsdFileCfg`` / ``Object.scale``. Returns: - AxisAlignedBoundingBox containing local min and max points. + OrientedBoundingBox containing local bounds. """ stage = Usd.Stage.Open(usd_path) if not stage: @@ -199,63 +183,38 @@ def compute_local_bounding_box_from_usd( bbox = compute_local_bounding_box_from_prim(stage, default_prim.GetPath().pathString) - usd_scale = _read_default_prim_scale(default_prim) - assert not any( - s == 0.0 for s in usd_scale - ), f"Default prim {default_prim.GetPath().pathString} has scale {usd_scale}" - composed_scale = (scale[0] / usd_scale[0], scale[1] / usd_scale[1], scale[2] / usd_scale[2]) - bbox = bbox.scaled(composed_scale) - return bbox + scale_array = np.asarray(scale, dtype=np.float32) + return OrientedBoundingBox( + center=bbox.center * bbox.center.new_tensor(scale_array), + half_extents=bbox.half_extents * bbox.half_extents.new_tensor(np.abs(scale_array)), + rotation_xyzw=bbox.rotation_xyzw, + ) def compute_local_bounding_box_from_prim( stage: Usd.Stage, prim_path: str, -) -> AxisAlignedBoundingBox: - """Compute the local bounding box of a specific prim (relative to prim's transform origin). +) -> OrientedBoundingBox: + """Compute axis-aligned geometry in the prim's local frame. Args: stage: The USD stage containing the prim. prim_path: Path to the prim to compute the bounding box for. Returns: - AxisAlignedBoundingBox containing the local min and max points relative to the - prim's own origin. - - Raises: - ValueError: If the prim is not found at the given path. + Axis-aligned bounds in the prim's local frame. """ prim = stage.GetPrimAtPath(prim_path) if not prim: raise ValueError(f"No prim found at path {prim_path}") - # Compute the world-space bounding box of the prim bbox_cache = UsdGeom.BBoxCache(Usd.TimeCode.Default(), includedPurposes=[UsdGeom.Tokens.default_]) - bbox = bbox_cache.ComputeWorldBound(prim) + bbox = bbox_cache.ComputeUntransformedBound(prim) bbox_range = bbox.ComputeAlignedBox() + local_min = bbox_range.GetMin() + local_max = bbox_range.GetMax() - # Get world-space min/max - world_min = bbox_range.GetMin() - world_max = bbox_range.GetMax() - - # Get the target prim's world position to compute local bounding box - prim_xformable = UsdGeom.Xformable(prim) - prim_world_transform = prim_xformable.ComputeLocalToWorldTransform(Usd.TimeCode.Default()) - prim_world_pos = prim_world_transform.ExtractTranslation() - - # Compute local bounding box by subtracting the prim's own world position - local_min = Gf.Vec3d( - world_min[0] - prim_world_pos[0], - world_min[1] - prim_world_pos[1], - world_min[2] - prim_world_pos[2], - ) - local_max = Gf.Vec3d( - world_max[0] - prim_world_pos[0], - world_max[1] - prim_world_pos[1], - world_max[2] - prim_world_pos[2], - ) - - return AxisAlignedBoundingBox( + return OrientedBoundingBox.from_min_max( min_point=(local_min[0], local_min[1], local_min[2]), max_point=(local_max[0], local_max[1], local_max[2]), ) @@ -265,75 +224,22 @@ def extract_trimesh_from_usd( usd_path: str, scale: tuple[float, float, float] = (1.0, 1.0, 1.0), ) -> trimesh.Trimesh: - """Extract all UsdGeom.Mesh prims from a USD into a single trimesh. + """Extract mesh geometry under a USD's default prim into a single trimesh. - Scale is applied per-vertex in local frame before the prim-to-world transform. + This is the public alias for ``extract_trimesh_from_usd_path``. + Unlike the legacy whole-stage extractor, it returns default-prim-local geometry + and applies scale after child transforms so mesh and bounding-box frames agree. All scale components must be positive (negative flips winding/SDF sign). - Other Gprim geometry is rejected, not silently dropped. + Other Gprim geometry under the default prim is rejected, not silently dropped. Args: usd_path: Path to the .usd/.usda/.usdc file. - scale: (sx, sy, sz) per-axis scale factors applied in local frame. + scale: (sx, sy, sz) per-axis scale factors applied in the default-prim frame. Returns: - Combined trimesh with per-prim world transforms baked in. + Combined trimesh in the scaled default-prim frame. """ - assert all( - s > 0 for s in scale - ), f"All scale components must be positive (negative scale flips winding/SDF sign), got {scale}" - - stage = Usd.Stage.Open(usd_path) - if stage is None: - raise ValueError(f"Failed to open USD: {usd_path}") - - all_verts: list[np.ndarray] = [] - all_faces: list[list[int]] = [] - skipped_gprims: list[str] = [] - offset = 0 - - for prim in stage.Traverse(): - if not prim.IsA(UsdGeom.Mesh): - if prim.IsA(UsdGeom.Gprim): - skipped_gprims.append(str(prim.GetPath())) - continue - mesh_prim = UsdGeom.Mesh(prim) - points = mesh_prim.GetPointsAttr().Get() - face_vertex_counts = mesh_prim.GetFaceVertexCountsAttr().Get() - face_vertex_indices = mesh_prim.GetFaceVertexIndicesAttr().Get() - if points is None or face_vertex_counts is None or face_vertex_indices is None: - continue - - xform = UsdGeom.Xformable(prim) - world_tf = np.array(xform.ComputeLocalToWorldTransform(Usd.TimeCode.Default())) - - verts = np.asarray(points, dtype=np.float64) - verts_scaled = verts * np.array(scale, dtype=np.float64) - verts_h = np.hstack([verts_scaled, np.ones((len(verts_scaled), 1))]) - verts_world = (verts_h @ world_tf)[:, :3] - - # Fan-triangulate faces - idx = 0 - for count in face_vertex_counts: - for k in range(1, count - 1): - all_faces.append([ - face_vertex_indices[idx] + offset, - face_vertex_indices[idx + k] + offset, - face_vertex_indices[idx + k + 1] + offset, - ]) - idx += count - - all_verts.append(verts_world) - offset += len(verts_world) - - if all_verts: - if skipped_gprims: - print(f"Unsupported non-mesh geometry in {usd_path}: {', '.join(skipped_gprims)}") - return trimesh.Trimesh(vertices=np.vstack(all_verts), faces=np.array(all_faces, dtype=np.int32)) - if skipped_gprims: - raise UnsupportedCollisionGeometryError( - f"Unsupported non-mesh geometry in {usd_path}: {', '.join(skipped_gprims)}" - ) - raise NoCollisionMeshError(f"No mesh geometry found in {usd_path}") + return extract_trimesh_from_usd_path(usd_path, scale) def extract_trimesh_from_prim( diff --git a/isaaclab_arena/utils/usd_pose_helpers.py b/isaaclab_arena/utils/usd_pose_helpers.py index 16e92908fb..295f7f0e05 100644 --- a/isaaclab_arena/utils/usd_pose_helpers.py +++ b/isaaclab_arena/utils/usd_pose_helpers.py @@ -37,7 +37,11 @@ def get_prim_pose_in_default_prim_frame(prim: Usd.Prim, stage: Usd.Stage) -> Pos default_T_world = default_T_world.GetInverse() prim_T_default = prim_T_world * default_T_world - pos, rot, _ = UsdSkel.DecomposeTransform(prim_T_default) + pos, rot, scale = UsdSkel.DecomposeTransform(prim_T_default) + assert all(abs(component - 1.0) < 1e-6 for component in scale), ( + "Referenced prim transform relative to the default prim must have unit scale; " + "Pose and oriented bounding boxes only support rigid transforms." + ) rot_tuple = (rot.GetImaginary()[0], rot.GetImaginary()[1], rot.GetImaginary()[2], rot.GetReal()) pos_tuple = (pos[0], pos[1], pos[2]) return Pose(position_xyz=pos_tuple, rotation_xyzw=rot_tuple) diff --git a/isaaclab_arena/utils/yaw.py b/isaaclab_arena/utils/yaw.py index df08db425d..aed27f97a5 100644 --- a/isaaclab_arena/utils/yaw.py +++ b/isaaclab_arena/utils/yaw.py @@ -54,15 +54,15 @@ def yaw_from_quat_xyzw(quat_xyzw: tuple[float, float, float, float]) -> float: def rotate_quat_by_yaw( base_xyzw: tuple[float, float, float, float], yaw_rad: float ) -> tuple[float, float, float, float]: - """Rotate base_xyzw (xyzw) by an extra yaw about Z. Returns base unchanged when yaw is 0.""" + """Apply an extra world-Z yaw as ``yaw_quat ⊗ base_xyzw``.""" yaw_rad = wrap_angle_to_pi(yaw_rad) # keep half-angle small for precision; canonicalize 2pi -> 0 if yaw_rad == 0.0: return base_xyzw bx, by, bz, bw = base_xyzw sz = math.sin(yaw_rad / 2.0) cz = math.cos(yaw_rad / 2.0) - # Hamilton product base ⊗ (0, 0, sz, cz). Both rotations are about Z, so they commute. - return (bx * cz + by * sz, -bx * sz + by * cz, bz * cz + bw * sz, -bz * sz + bw * cz) + # Hamilton product (0, 0, sz, cz) ⊗ base. + return (bx * cz - by * sz, bx * sz + by * cz, bz * cz + bw * sz, -bz * sz + bw * cz) def rotate_points_by_yaw(points: torch.Tensor, yaw: float) -> torch.Tensor: diff --git a/isaaclab_arena_curobo/ik_reachability_validator.py b/isaaclab_arena_curobo/ik_reachability_validator.py index 6bb779b02c..51177239f1 100644 --- a/isaaclab_arena_curobo/ik_reachability_validator.py +++ b/isaaclab_arena_curobo/ik_reachability_validator.py @@ -14,40 +14,47 @@ import torch from typing import TYPE_CHECKING -from isaaclab_arena.relations.placement_events import get_base_rotation_per_asset from isaaclab_arena.relations.placement_validation import PlacementCheck from isaaclab_arena.relations.placement_validator_registry import register_validator from isaaclab_arena.relations.placement_validators import PlacementValidator from isaaclab_arena.relations.relations import RequiresReachability, get_anchor_objects from isaaclab_arena.utils.pose import Pose -from isaaclab_arena.utils.yaw import rotate_quat_by_yaw, yaw_from_quat_xyzw from isaaclab_arena_curobo.embodiment_curobo_registry import get_embodiment_curobo_cfg from isaaclab_arena_curobo.ik_solver import CuroboIKSolver from isaaclab_arena_curobo.utils.frame_utils import top_down_grasp_pose_from_world_poses -from isaaclab_arena_curobo.utils.ik_solver_utils import get_aabb_collision_cuboid_for_object, solve_ik_feasibility +from isaaclab_arena_curobo.utils.ik_solver_utils import get_obb_collision_cuboid_for_object, solve_ik_feasibility if TYPE_CHECKING: from isaaclab_arena.assets.object_base import ObjectBase from isaaclab_arena.relations.collision_object import CollisionObject from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox def get_object_world_pose_from_layout( positions: dict[ObjectBase, tuple[float, float, float]], - orientations: dict[ObjectBase, float], + rotations: dict[ObjectBase, tuple[float, float, float, float]], obj: ObjectBase, - base_rotations: dict, + anchors: set[ObjectBase], ) -> Pose: """Return the world pose an object gets under a layout.""" - pos_w = positions[obj] - base_quat_xyzw = base_rotations[obj] - marker_yaw = yaw_from_quat_xyzw(base_quat_xyzw) - total_yaw = orientations.get(obj, marker_yaw) - quat_w_xyzw = rotate_quat_by_yaw(base_quat_xyzw, total_yaw - marker_yaw) + if obj in rotations: + rotation_xyzw = rotations[obj] + elif obj in anchors: + fixed_pose = obj.get_initial_pose() + assert isinstance(fixed_pose, Pose), f"Anchor '{obj.name}' must have a fixed Pose." + rotation_xyzw = fixed_pose.rotation_xyzw + else: + rotation_xyzw = (0.0, 0.0, 0.0, 1.0) + rotation = torch.tensor(rotation_xyzw, dtype=torch.float32) + assert rotation.shape == (4,), f"Rotation for '{obj.name}' must have shape (4,)." + assert torch.isfinite(rotation).all(), f"Rotation for '{obj.name}' must be finite." + assert torch.isclose( + torch.linalg.vector_norm(rotation), torch.tensor(1.0), atol=1e-5, rtol=1e-5 + ), f"Rotation for '{obj.name}' must be a unit quaternion." return Pose( - position_xyz=tuple(float(v) for v in pos_w), - rotation_xyzw=tuple(float(v) for v in quat_w_xyzw), + position_xyz=tuple(float(v) for v in positions[obj]), + rotation_xyzw=tuple(float(v) for v in rotation_xyzw), ) @@ -95,16 +102,17 @@ def is_available(cls, params: ObjectPlacerParams) -> bool: def validate_batch( self, positions: list[dict[ObjectBase, tuple[float, float, float]]], - orientations: list[dict[ObjectBase, float]], - bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]], + rotations: list[dict[ObjectBase, tuple[float, float, float, float]]], + bboxes: list[dict[ObjectBase, OrientedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: - return [self._validate(positions[i], orientations[i]) for i in range(len(positions))] + return [self._validate(positions[i], rotations[i], bboxes[i]) for i in range(len(positions))] def _validate( self, positions: dict[ObjectBase, tuple[float, float, float]], - orientations: dict[ObjectBase, float], + rotations: dict[ObjectBase, tuple[float, float, float, float]], + bboxes: dict[ObjectBase, OrientedBoundingBox], ) -> bool: """Whether the robot can reach a top-down grasp at the target objects in one candidate layout. @@ -114,13 +122,15 @@ def _validate( """ objects = list(positions.keys()) anchors = set(get_anchor_objects(objects)) - base_rotations = get_base_rotation_per_asset(objects) - world_poses = { - obj: get_object_world_pose_from_layout(positions, orientations, obj, base_rotations) for obj in objects - } + world_poses = {obj: get_object_world_pose_from_layout(positions, rotations, obj, anchors) for obj in objects} cuboids = [ - get_aabb_collision_cuboid_for_object(obj, world_poses[obj].position_xyz, world_poses[obj].rotation_xyzw) + get_obb_collision_cuboid_for_object( + obj, + bboxes[obj], + world_poses[obj].position_xyz, + world_poses[obj].rotation_xyzw, + ) for obj in objects ] self._solver.update_world(cuboids, self._base_pos, self._base_quat_xyzw) diff --git a/isaaclab_arena_curobo/ik_solver.py b/isaaclab_arena_curobo/ik_solver.py index ef8f327b0a..131a79f31d 100644 --- a/isaaclab_arena_curobo/ik_solver.py +++ b/isaaclab_arena_curobo/ik_solver.py @@ -19,7 +19,7 @@ from isaaclab_arena.utils.device import resolve_cuda_device from isaaclab_arena_curobo.curobo_embodiment_cfg import CuroboEmbodimentCfg from isaaclab_arena_curobo.embodiment_curobo_registry import get_embodiment_curobo_cfg -from isaaclab_arena_curobo.utils.ik_solver_utils import AABBCollisionCuboid, world_config_from_cuboids +from isaaclab_arena_curobo.utils.ik_solver_utils import OrientedCollisionCuboid, world_config_from_cuboids from isaaclab_arena_curobo.utils.robot_cfg_utils import load_patched_robot_yaml @@ -137,7 +137,7 @@ def _make_pose( def update_world( self, - cuboids: list[AABBCollisionCuboid], + cuboids: list[OrientedCollisionCuboid], robot_base_pos_w: tuple[float, float, float], robot_base_quat_w_xyzw: tuple[float, float, float, float], ) -> None: diff --git a/isaaclab_arena_curobo/tests/test_ik_reachability_validator.py b/isaaclab_arena_curobo/tests/test_ik_reachability_validator.py index 76eb0fceed..86ad42e8b6 100644 --- a/isaaclab_arena_curobo/tests/test_ik_reachability_validator.py +++ b/isaaclab_arena_curobo/tests/test_ik_reachability_validator.py @@ -6,9 +6,9 @@ """Logic tests for ReachabilityValidator, with cuRobo mocked out. Exercises what the validator does around cuRobo -- reconstruct object poses from a layout, build one -collision cuboid per object, and IK-check one grasp per movable (non-anchor) object -- against a real -geometry-solved layout, asserting the per-layout ``validate_batch`` verdict. The cuRobo solver build -and the batched IK solve are patched, so no GPU or cuRobo install is needed; the pure-math grasp +collision cuboid per object, and IK-check one grasp per task-marked reachability target -- against a +real geometry-solved layout, asserting the per-layout ``validate_batch`` verdict. The cuRobo solver +build and batched IK solve are patched, so no GPU or cuRobo install is needed; pure-math grasp reconstruction runs for real on CPU. """ @@ -20,25 +20,29 @@ import pytest -def _make_desk_box_pool(num_envs: int = 1, min_layouts_per_env: int = 2): +def _make_desk_box_pool( + num_envs: int = 1, + min_layouts_per_env: int = 2, + desk_rotation: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0), +): """Build a small valid desk (anchor) + box (On desk) pool and return it.""" from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On, RequiresReachability from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), ) - desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=desk_rotation)) desk.add_relation(IsAnchor()) box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=0.01)) box.add_relation(RequiresReachability()) @@ -85,8 +89,19 @@ def _fake_ik(solver, target_poses, **kwargs): captured["num_grasps"] = num return feasible, torch.zeros(num), torch.zeros(num) + def _fake_cuboid(obj, bbox, position, rotation): + captured.setdefault("cuboid_bboxes", {})[obj.name] = bbox + captured.setdefault("cuboid_poses", {})[obj.name] = (position, rotation) + return obj.name + + def _fake_grasp(position, rotation, *args, **kwargs): + captured.setdefault("grasp_poses", []).append((position, rotation)) + return torch.tensor((*position, *rotation), dtype=torch.float32) + monkeypatch.setattr(mod, "CuroboIKSolver", _make_solver) monkeypatch.setattr(mod, "solve_ik_feasibility", _fake_ik) + monkeypatch.setattr(mod, "get_obb_collision_cuboid_for_object", _fake_cuboid) + monkeypatch.setattr(mod, "top_down_grasp_pose_from_world_poses", _fake_grasp) monkeypatch.setattr(mod, "get_embodiment_curobo_cfg", lambda embodiment: None) return captured @@ -107,12 +122,12 @@ def _make_two_box_pool(num_envs: int = 1, min_layouts_per_env: int = 2): from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On, RequiresReachability from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) @@ -120,7 +135,7 @@ def _make_two_box_pool(num_envs: int = 1, min_layouts_per_env: int = 2): for box_name in ("box_a", "box_b"): box = DummyObject( name=box_name, - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=0.01)) # Only box_a carries the RequiresReachability marker, so only it should be IK-checked. @@ -149,18 +164,18 @@ def _make_unstamped_desk_box_pool(num_envs: int = 1, min_layouts_per_env: int = from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On from isaaclab_arena.tests.dummy_object import DummyObject - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose desk = DummyObject( name="desk", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), ) desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) box = DummyObject( name="box", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), ) box.add_relation(On(desk, clearance_m=0.01)) @@ -188,27 +203,74 @@ def _make_reachability_validator(embodiment): return ReachabilityValidator(params) +def _layout_bboxes(layout): + """Return the complete candidate bounding-box map for a solved layout.""" + return {obj: obj.get_bounding_box() for obj in layout.positions} + + +@pytest.mark.curobo_deps +def test_collision_cuboid_composes_intrinsic_bbox_rotation(): + """The cuRobo cuboid pose includes the OBB's asset-local orientation.""" + from isaaclab_arena.tests.dummy_object import DummyObject + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox + from isaaclab_arena_curobo.utils.ik_solver_utils import get_obb_collision_cuboid_for_object + + bbox_rotation = (0.0, 0.0, 2**-0.5, 2**-0.5) + obj = DummyObject( + name="rotated_box", + bounding_box=OrientedBoundingBox( + center=(0.25, 0.0, 0.0), + half_extents=(0.2, 0.1, 0.05), + rotation_xyzw=bbox_rotation, + ), + ) + + cuboid = get_obb_collision_cuboid_for_object( + obj, + obj.get_bounding_box(), + pos_w=(1.0, 2.0, 3.0), + quat_w_xyzw=(0.0, 0.0, 0.0, 1.0), + ) + + assert cuboid.pose_W_O.position_xyz == pytest.approx((1.25, 2.0, 3.0)) + assert cuboid.pose_W_O.rotation_xyzw == pytest.approx(bbox_rotation) + + @pytest.mark.curobo_deps def test_validator_accepts_when_all_grasps_feasible(monkeypatch): - """A layout passes when every movable-object grasp is feasible.""" + """Full world quaternions reach cuboid and grasp construction unchanged.""" captured = _patch_curobo(monkeypatch, feasible_fn=lambda n: [True] * n) validator = _make_reachability_validator(_fake_embodiment()) - layout = _make_desk_box_pool().layouts_per_env()[0][0] - assert validator.validate_batch([layout.positions], [layout.orientations], [{}], []) == [True] + anchor_rotation = (0.5, 0.0, 0.0, 3**0.5 / 2) + movable_rotation = (0.5, 0.5, 0.5, 0.5) + layout = _make_desk_box_pool(desk_rotation=anchor_rotation).layouts_per_env()[0][0] + box = next(obj for obj in layout.positions if obj.name == "box") + layout.rotations[box] = movable_rotation + assert validator.validate_batch([layout.positions], [layout.rotations], [_layout_bboxes(layout)], []) == [True] # One collision cuboid per object (desk + box); one grasp per movable object (box only, desk is anchor). assert len(captured["solver"].world_cuboids) == 2 assert captured["num_grasps"] == 1 + assert captured["cuboid_poses"]["desk"][1] == pytest.approx(anchor_rotation) + box_cuboid_position, box_cuboid_rotation = captured["cuboid_poses"]["box"] + assert box_cuboid_position == pytest.approx(layout.positions[box]) + assert box_cuboid_rotation == pytest.approx(movable_rotation) + assert len(captured["grasp_poses"]) == 1 + grasp_position, grasp_rotation = captured["grasp_poses"][0] + assert grasp_position == pytest.approx(layout.positions[box]) + assert grasp_rotation == pytest.approx(movable_rotation) @pytest.mark.curobo_deps def test_validator_rejects_when_any_grasp_infeasible(monkeypatch): - """A layout fails when any movable-object grasp is infeasible.""" - _patch_curobo(monkeypatch, feasible_fn=lambda n: [False] * n) + """Omitted movable rotations use identity before an infeasible grasp rejects the layout.""" + captured = _patch_curobo(monkeypatch, feasible_fn=lambda n: [False] * n) validator = _make_reachability_validator(_fake_embodiment()) layout = _make_desk_box_pool().layouts_per_env()[0][0] - assert validator.validate_batch([layout.positions], [layout.orientations], [{}], []) == [False] + layout.rotations.clear() + assert validator.validate_batch([layout.positions], [layout.rotations], [_layout_bboxes(layout)], []) == [False] + assert captured["cuboid_poses"]["box"][1] == pytest.approx((0.0, 0.0, 0.0, 1.0)) @pytest.mark.curobo_deps @@ -218,7 +280,7 @@ def test_validator_checks_only_stamped_objects(monkeypatch): validator = _make_reachability_validator(_fake_embodiment()) layout = _make_two_box_pool().layouts_per_env()[0][0] - assert validator.validate_batch([layout.positions], [layout.orientations], [{}], []) == [True] + assert validator.validate_batch([layout.positions], [layout.rotations], [_layout_bboxes(layout)], []) == [True] # Two movable boxes exist, but only the stamped one (box_a) contributes a grasp. assert captured["num_grasps"] == 1 @@ -231,10 +293,35 @@ def test_validator_passes_trivially_and_warns_when_no_targets(monkeypatch, capsy layout = _make_unstamped_desk_box_pool().layouts_per_env()[0][0] # Two layouts through the same validator: the warning must print once, not once per candidate. + bboxes = _layout_bboxes(layout) assert validator.validate_batch( - [layout.positions, layout.positions], [layout.orientations, layout.orientations], [{}, {}], [] + [layout.positions, layout.positions], + [layout.rotations, layout.rotations], + [bboxes, bboxes], + [], ) == [True, True] # No grasp was ever solved (the IK path is skipped entirely when there are no targets). assert "num_grasps" not in captured assert capsys.readouterr().out.count("resolved zero reachability targets") == 1 + + +@pytest.mark.curobo_deps +def test_validator_passes_candidate_bbox_to_cuboid_construction(monkeypatch): + """The candidate's heterogeneous box, not the object's default box, defines its cuboid.""" + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox + + captured = _patch_curobo(monkeypatch, feasible_fn=lambda n: [True] * n) + validator = _make_reachability_validator(_fake_embodiment()) + layout = _make_desk_box_pool().layouts_per_env()[0][0] + box = next(obj for obj in layout.positions if obj.name == "box") + candidate_bbox = OrientedBoundingBox( + center=(0.4, -0.2, 0.1), + half_extents=(0.3, 0.2, 0.1), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), + ) + bboxes = _layout_bboxes(layout) + bboxes[box] = candidate_bbox + + assert validator.validate_batch([layout.positions], [layout.rotations], [bboxes], []) == [True] + assert captured["cuboid_bboxes"]["box"] is candidate_bbox diff --git a/isaaclab_arena_curobo/utils/ik_solver_utils.py b/isaaclab_arena_curobo/utils/ik_solver_utils.py index 9a516444e0..69fa02993d 100644 --- a/isaaclab_arena_curobo/utils/ik_solver_utils.py +++ b/isaaclab_arena_curobo/utils/ik_solver_utils.py @@ -23,6 +23,7 @@ from curobo.wrap.reacher.ik_solver import IKSolver from isaaclab_mimic.motion_planners.curobo.curobo_planner import CuroboPlanner + from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena_curobo.ik_solver import CuroboIKSolver @@ -39,8 +40,8 @@ def resolve_ik_solver(ik_solver_context: CuroboIKSolver | CuroboPlanner) -> IKSo @dataclass -class AABBCollisionCuboid: - """A collision obstacle described by an axis-aligned bounding box in the world frame. +class OrientedCollisionCuboid: + """A collision obstacle described by an oriented box in the world frame. ``dims_xyz`` are full extents (edge lengths), matching cuRobo's ``Cuboid.dims``. """ @@ -50,32 +51,33 @@ class AABBCollisionCuboid: pose_W_O: Pose = field(default_factory=Pose.identity) -def get_aabb_collision_cuboid_for_object( - obj: ObjectBase, pos_w: tuple[float, float, float], quat_w_xyzw: tuple[float, ...] -) -> AABBCollisionCuboid: - """Axis-aligned bounding-box collision cuboid for an object at its layout pose (world frame). +def get_obb_collision_cuboid_for_object( + obj: ObjectBase, + bbox: OrientedBoundingBox, + pos_w: tuple[float, float, float], + quat_w_xyzw: tuple[float, ...], +) -> OrientedCollisionCuboid: + """Oriented collision cuboid for an object at its layout pose. - The bounding box is object-local, so its center offset is rotated by the object's world orientation - and added to the root position -- placing e.g. a table box at its true mid-height rather than at the - root. + The box center and orientation are transformed from object-local to world coordinates. """ - bbox = obj.get_bounding_box() - dims = tuple(float(v) for v in bbox.size[0].tolist()) + dims = tuple(float(v) for v in (2.0 * bbox.half_extents[0]).tolist()) quat_t = torch.tensor(quat_w_xyzw, dtype=torch.float32) rotation = math_utils.matrix_from_quat(quat_t.unsqueeze(0))[0] center_world = torch.tensor(pos_w, dtype=torch.float32) + rotation @ bbox.center[0].to(torch.float32) - return AABBCollisionCuboid( + box_quat_world = math_utils.quat_mul(quat_t.unsqueeze(0), bbox.rotation_xyzw.to(torch.float32))[0] + return OrientedCollisionCuboid( name=obj.name, dims_xyz=dims, pose_W_O=Pose( position_xyz=tuple(float(v) for v in center_world.tolist()), - rotation_xyzw=tuple(float(v) for v in quat_w_xyzw), + rotation_xyzw=tuple(float(v) for v in box_quat_world.tolist()), ), ) def world_config_from_cuboids( - cuboids: list[AABBCollisionCuboid], + cuboids: list[OrientedCollisionCuboid], robot_base_pos_w: tuple[float, float, float], robot_base_quat_w_xyzw: tuple[float, float, float, float], device: str | torch.device | None = None, diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/simapp/asset_usd.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/simapp/asset_usd.py index eb7bcc8b07..a4db014746 100644 --- a/isaaclab_arena_examples/agentic_environment_generation/review_gui/simapp/asset_usd.py +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/simapp/asset_usd.py @@ -20,7 +20,8 @@ def aabb_dimensions_from_asset(asset: ObjectBase) -> AabbDimensionsM | None: """Return local axis-aligned bounding box size (x, y, z) in meters for one live asset.""" try: bbox = asset.get_bounding_box() - size = bbox.size[0] + min_point, max_point = bbox.get_axis_aligned_bounds() + size = (max_point - min_point)[0] return (float(size[0]), float(size[1]), float(size[2])) except Exception as exc: name = getattr(asset, "name", "?") diff --git a/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py b/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py index 398f916d4c..1607ac904b 100644 --- a/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py +++ b/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py @@ -15,7 +15,7 @@ from isaaclab_arena.relations.relation_solver import RelationSolver from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, NextTo, On, Side, get_anchor_objects -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose from isaaclab_arena_examples.relations.example_object import ExampleObject from isaaclab_arena_examples.relations.relation_solver_visualizer import RelationSolverVisualizer @@ -26,31 +26,37 @@ def run_dummy_object_placer_demo(): """Run the ObjectPlacer demo with dummy objects and a single anchor.""" # Create objects with bounding boxes desk = ExampleObject( - name="desk", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)) + name="desk", bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)) ) # Central object on the desk center_box = ExampleObject( - name="center_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.15)) + name="center_box", + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.15)), ) # Objects placed on each side of center_box right_box = ExampleObject( - name="right_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.1)) + name="right_box", + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.1)), ) left_box = ExampleObject( - name="left_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.1)) + name="left_box", + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.1)), ) front_box = ExampleObject( - name="front_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.08)) + name="front_box", + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.08)), ) back_box = ExampleObject( - name="back_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.08)) + name="back_box", + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.08)), ) # Box on top of center_box top_box = ExampleObject( - name="top_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.08)) + name="top_box", + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.08)), ) # Mark desk as the anchor for relation solving (not subject to optimization) @@ -107,23 +113,23 @@ def run_dummy_multi_anchor_demo(): # Create anchor objects (fixed positions) table = ExampleObject( name="table", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 0.6, 0.75)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 0.6, 0.75)), ) chair = ExampleObject( name="chair", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.5, 0.5, 0.45)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.5, 0.5, 0.45)), ) mug = ExampleObject( name="mug", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.1)), ) book = ExampleObject( name="book", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.15, 0.03)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.15, 0.03)), ) bin_obj = ExampleObject( name="bin", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.4)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.4)), ) # Anchor objects (fixed positions) @@ -169,22 +175,22 @@ def run_dummy_no_collision_demo(): # Create table (anchor) and three boxes table = ExampleObject( name="table", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.8, 0.6, 0.4)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.8, 0.6, 0.4)), ) table.add_relation(IsAnchor()) table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) box_a = ExampleObject( name="box_a", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.1)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.1)), ) box_b = ExampleObject( name="box_b", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.08)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.08)), ) box_c = ExampleObject( name="box_c", - bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.18, 0.1, 0.06)), + bounding_box=OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.18, 0.1, 0.06)), ) # Boxes on table; no-overlap is handled automatically by the solver diff --git a/isaaclab_arena_examples/relations/example_object.py b/isaaclab_arena_examples/relations/example_object.py index e4354c3551..fc78239988 100644 --- a/isaaclab_arena_examples/relations/example_object.py +++ b/isaaclab_arena_examples/relations/example_object.py @@ -5,15 +5,15 @@ from __future__ import annotations from isaaclab_arena.relations.placement_asset import PlaceableAsset -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox class ExampleObject(PlaceableAsset): """Box-shaped placement asset for the relation-solver example notebooks, with no Isaac Sim dependency.""" - def __init__(self, name: str, bounding_box: AxisAlignedBoundingBox): + def __init__(self, name: str, bounding_box: OrientedBoundingBox): super().__init__(name=name) self._bounding_box = bounding_box - def get_bounding_box(self) -> AxisAlignedBoundingBox: + def get_bounding_box(self) -> OrientedBoundingBox: return self._bounding_box diff --git a/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py b/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py index 3c25634f59..a429d0e941 100644 --- a/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py +++ b/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py @@ -21,7 +21,7 @@ from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relation_solver_state import RelationSolverState from isaaclab_arena.relations.relations import IsAnchor, NextTo, Side -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox from isaaclab_arena.utils.pose import Pose from isaaclab_arena_examples.relations.example_object import ExampleObject @@ -93,34 +93,42 @@ def plot_loss_heatmap(X, Y, losses, parent, child, side, distance_m): parent_pose = parent.get_initial_pose() parent_bbox = parent.get_bounding_box() px, py, pz = parent_pose.position_xyz - pw, pd, ph = parent_bbox.size[0].tolist() + parent_min, parent_max = parent_bbox.get_axis_aligned_bounds() + pmin_x, pmin_y, _ = parent_min[0].tolist() + pmax_x, pmax_y, _ = parent_max[0].tolist() + pw, pd = pmax_x - pmin_x, pmax_y - pmin_y # Draw parent bounding box parent_rect = Rectangle( - (px - pw / 2, py - pd / 2), pw, pd, linewidth=3, edgecolor="blue", facecolor="none", label="Parent Object" + (px + pmin_x, py + pmin_y), pw, pd, linewidth=3, edgecolor="blue", facecolor="none", label="Parent Object" ) ax.add_patch(parent_rect) ax.plot(px, py, "b*", markersize=15, label="Parent Center") # Get child bounding box child_bbox = child.get_bounding_box() - cw, cd, ch = child_bbox.size[0].tolist() + child_min, child_max = child_bbox.get_axis_aligned_bounds() + cmin_x, cmin_y, _ = child_min[0].tolist() + cmax_x, cmax_y, _ = child_max[0].tolist() + cw, cd = cmax_x - cmin_x, cmax_y - cmin_y + aligned_x = px + (pmin_x + pmax_x - cmin_x - cmax_x) / 2 + aligned_y = py + (pmin_y + pmax_y - cmin_y - cmax_y) / 2 # Mark ideal position if side == Side.POSITIVE_X: - ideal_x, ideal_y = px + pw / 2 + distance_m + cw / 2, py + ideal_x, ideal_y = px + pmax_x + distance_m - cmin_x, aligned_y elif side == Side.NEGATIVE_X: - ideal_x, ideal_y = px - pw / 2 - distance_m - cw / 2, py + ideal_x, ideal_y = px + pmin_x - distance_m - cmax_x, aligned_y elif side == Side.NEGATIVE_Y: - ideal_x, ideal_y = px, py - pd / 2 - distance_m - cd / 2 + ideal_x, ideal_y = aligned_x, py + pmin_y - distance_m - cmax_y elif side == Side.POSITIVE_Y: - ideal_x, ideal_y = px, py + pd / 2 + distance_m + cd / 2 + ideal_x, ideal_y = aligned_x, py + pmax_y + distance_m - cmin_y ax.plot(ideal_x, ideal_y, "g*", markersize=15, label="Ideal Position") # Draw child bounding box at ideal position child_rect = Rectangle( - (ideal_x - cw / 2, ideal_y - cd / 2), + (ideal_x + cmin_x, ideal_y + cmin_y), cw, cd, linewidth=2, @@ -148,9 +156,9 @@ def plot_loss_heatmap(X, Y, losses, parent, child, side, distance_m): # %% def run_visualization_demo(): """Run the full visualization demo.""" - parent_bbox = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.5, 0.5, 0.1)) + parent_bbox = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.5, 0.5, 0.1)) parent_pos = (0.0, 0.0, 0.05) - child_bbox = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.15)) + child_bbox = OrientedBoundingBox.from_min_max(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.15)) distance_m = 0.1 # Create parent object diff --git a/isaaclab_arena_examples/relations/relation_solver_visualizer.py b/isaaclab_arena_examples/relations/relation_solver_visualizer.py index 4ac92562ab..d8497876fd 100644 --- a/isaaclab_arena_examples/relations/relation_solver_visualizer.py +++ b/isaaclab_arena_examples/relations/relation_solver_visualizer.py @@ -11,7 +11,7 @@ import plotly.graph_objects as go -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.bounding_box import OrientedBoundingBox if TYPE_CHECKING: from isaaclab_arena.assets.object import Object @@ -69,7 +69,7 @@ def _get_color(self, idx: int) -> str: def _create_wireframe_box( self, - bbox: AxisAlignedBoundingBox, + bbox: OrientedBoundingBox, position: tuple[float, float, float], color: str, name: str, @@ -89,29 +89,9 @@ def _create_wireframe_box( Returns: Scatter3d trace forming the wireframe """ - # Compute world-space corners: world = position + local offset - # This matches how the loss strategies compute world extents x, y, z = position - local_min = bbox.min_point[0].tolist() - local_max = bbox.max_point[0].tolist() - x_min = x + local_min[0] - x_max = x + local_max[0] - y_min = y + local_min[1] - y_max = y + local_max[1] - z_min = z + local_min[2] - z_max = z + local_max[2] - - # 8 corners of the box (same ordering as get_corners_at) - corners = [ - [x_min, y_min, z_min], # 0: Bottom-front-left - [x_max, y_min, z_min], # 1: Bottom-front-right - [x_max, y_max, z_min], # 2: Bottom-back-right - [x_min, y_max, z_min], # 3: Bottom-back-left - [x_min, y_min, z_max], # 4: Top-front-left - [x_max, y_min, z_max], # 5: Top-front-right - [x_max, y_max, z_max], # 6: Top-back-right - [x_min, y_max, z_max], # 7: Top-back-left - ] + corners = bbox.get_corners()[0].tolist() + corners = [[corner[0] + x, corner[1] + y, corner[2] + z] for corner in corners] # Define edges as pairs of corner indices (matching get_corners ordering) edges = [ @@ -338,7 +318,7 @@ def plot_loss_history(self) -> go.Figure: def _get_wireframe_coords( self, - bbox: AxisAlignedBoundingBox, + bbox: OrientedBoundingBox, position: tuple[float, float, float], ) -> tuple[list, list, list]: """Get wireframe coordinates for a bounding box at a position. @@ -350,27 +330,9 @@ def _get_wireframe_coords( Returns: Tuple of (x_coords, y_coords, z_coords) lists for the wireframe """ - # Compute world-space corners: world = position + local offset x, y, z = position - local_min = bbox.min_point[0].tolist() - local_max = bbox.max_point[0].tolist() - x_min = x + local_min[0] - x_max = x + local_max[0] - y_min = y + local_min[1] - y_max = y + local_max[1] - z_min = z + local_min[2] - z_max = z + local_max[2] - - corners = [ - [x_min, y_min, z_min], - [x_max, y_min, z_min], - [x_max, y_max, z_min], - [x_min, y_max, z_min], - [x_min, y_min, z_max], - [x_max, y_min, z_max], - [x_max, y_max, z_max], - [x_min, y_max, z_max], - ] + corners = bbox.get_corners()[0].tolist() + corners = [[corner[0] + x, corner[1] + y, corner[2] + z] for corner in corners] edges = [ (0, 1), @@ -420,10 +382,10 @@ def animate_optimization(self) -> go.Figure: for idx, obj in enumerate(self.objects): pos = positions[idx] bbox = obj.get_bounding_box() - half_size = (bbox.size[0] / 2).tolist() - all_x.extend([pos[0] - half_size[0], pos[0] + half_size[0]]) - all_y.extend([pos[1] - half_size[1], pos[1] + half_size[1]]) - all_z.extend([pos[2] - half_size[2], pos[2] + half_size[2]]) + local_min, local_max = bbox.get_axis_aligned_bounds() + all_x.extend([pos[0] + local_min[0, 0].item(), pos[0] + local_max[0, 0].item()]) + all_y.extend([pos[1] + local_min[0, 1].item(), pos[1] + local_max[0, 1].item()]) + all_z.extend([pos[2] + local_min[0, 2].item(), pos[2] + local_max[0, 2].item()]) padding = 0.1 x_range = [min(all_x) - padding, max(all_x) + padding]