Introduce OrientedBoundingBox

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

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

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

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

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
 

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)

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)."""