Support Multi-part Collision Geometry

isaaclab_arena/embodiments/droid/droid.py

@@ -8,7 +8,7 @@
 import torch
 import warnings
 from abc import ABC
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import isaaclab.envs.mdp as mdp_isaac_lab
 import isaaclab.sim as sim_utils
@@ -43,6 +43,12 @@
 from isaaclab_arena.utils.cameras import ArenaCameraCfg
 from isaaclab_arena.utils.pose import Pose, PosePerEnv, translate_by_xyz_offset
 
+if TYPE_CHECKING:
+    import trimesh
+
+    from isaaclab_arena.relations.collision_object import CollisionComponent
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
 # The base stand's x/y footprint.
 _STAND_FOOTPRINT_SCALE_XY: tuple[float, float] = (1.2, 1.2)
 # The default stand height.
@@ -71,6 +77,8 @@ def __init__(
     ):
         super().__init__(enable_cameras, initial_pose, concatenate_observation_terms, arm_mode)
         self.scene_config = DroidSceneCfg()
+        # Lazily-extracted stand collision mesh, cached so its USD is opened once.
+        self._stand_collision_mesh_cache = None
         # ``stand_height_m`` is an absolute height in meters; convert it to the z-scale the USD needs.
         stand_unit_height = _stand_unit_height_m(self.scene_config.stand.spawn.usd_path)
         self.scene_config.stand.spawn.scale = (*_STAND_FOOTPRINT_SCALE_XY, stand_height_m / stand_unit_height)
@@ -124,6 +132,64 @@ def _update_scene_cfg_with_robot_initial_pose(self, scene_config: Any, pose: Pos
         scene_config.stand.init_state.rot = pose.rotation_xyzw
         return scene_config
 
+    def get_relation_bounding_box(self) -> "AxisAlignedBoundingBox":
+        """Use the stand footprint for ``On``/``NextTo``, not the arm's wider envelope.
+
+        The stand is the part of the mobile base that sits against the support surface, so its
+        footprint is what a proximity/support relation should measure. The arm reaches out well
+        beyond it, and a union box padded to the arm's extent would place the base too far away.
+        """
+        return self._stand_bounding_box()
+
+    def get_collision_components(self) -> list["CollisionComponent"]:
+        """Expose the robot arm and the stand as two separate, identity-posed sub-volumes.
+
+        A single union AABB over the arm and the offset stand is a poor collision proxy: rotating it
+        sweeps large empty regions. Keeping them separate lets a collision check test each against
+        obstacles independently. The solved base pose is written to both the robot and stand prims
+        (see ``_update_scene_cfg_with_robot_initial_pose``), so both components share the Droid root
+        frame and need no local offset.
+        """
+        from isaaclab_arena.relations.collision_object import CollisionComponent
+
+        return [
+            CollisionComponent(
+                name="robot",
+                local_pose=Pose.identity(),
+                bounding_box=self.get_bounding_box(),
+                mesh=self.get_collision_mesh(),
+            ),
+            CollisionComponent(
+                name="stand",
+                local_pose=Pose.identity(),
+                bounding_box=self._stand_bounding_box(),
+                mesh=self._stand_collision_mesh(),
+            ),
+        ]
+
+    def _stand_bounding_box(self) -> "AxisAlignedBoundingBox":
+        """Return the stand's root-relative bounds from its scaled USD geometry."""
+        from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd
+
+        spawn = self.scene_config.stand.spawn
+        assert spawn.usd_path is not None, "scene_config.stand must use a USD spawn for placement"
+        scale = tuple(spawn.scale or (1.0, 1.0, 1.0))
+        return compute_local_bounding_box_from_usd(spawn.usd_path, scale)
+
+    def _stand_collision_mesh(self) -> "trimesh.Trimesh | None":
+        """Return the stand's collision mesh from its USD, or ``None`` if it has no mesh geometry."""
+        if self._stand_collision_mesh_cache is None:
+            from isaaclab_arena.utils.usd_helpers import NoCollisionMeshError, extract_trimesh_from_usd_path
+
+            spawn = self.scene_config.stand.spawn
+            assert spawn.usd_path is not None, "scene_config.stand must use a USD spawn for placement"
+            scale = tuple(spawn.scale or (1.0, 1.0, 1.0))
+            try:
+                self._stand_collision_mesh_cache = extract_trimesh_from_usd_path(spawn.usd_path, scale)
+            except NoCollisionMeshError:
+                self._stand_collision_mesh_cache = None
+        return self._stand_collision_mesh_cache
+
     def set_initial_joint_pose(self, initial_joint_pose: list[float]) -> None:
         self.event_config.init_franka_arm_pose.params["default_pose"] = initial_joint_pose
 

isaaclab_arena/relations/collision_object.py

@@ -8,13 +8,38 @@
 from __future__ import annotations
 
 import trimesh
+from dataclasses import dataclass
 from typing import Protocol
 
 from isaaclab_arena.relations.collision_mode import CollisionMode
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange
 
 
+@dataclass(frozen=True)
+class CollisionComponent:
+    """A rigid sub-volume of a placement asset, fixed relative to the asset root.
+
+    A simple asset has one component; a compound asset (e.g. a robot on a stand)
+    has several. Each component's geometry lives in its own frame. With ``N`` the
+    number of environments, ``bounding_box`` is per-env ((N, 3) min/max points),
+    whereas ``local_pose`` is a single transform shared across all N (``Pose`` is
+    tuple-backed, not batched).
+    """
+
+    name: str
+    """Identifier for this component within the owning asset."""
+
+    local_pose: Pose
+    """Transform from the asset root frame to this component's frame, shared across all N envs."""
+
+    bounding_box: AxisAlignedBoundingBox
+    """Axis-aligned component-frame extents, per env ((N, 3) min and max points)."""
+
+    mesh: trimesh.Trimesh | None = None
+    """Collision mesh in the component frame; None means use ``bounding_box``."""
+
+
 class CollisionObject(Protocol):
     """Object the collision solver can query for pose, bounds, and mesh."""
 

isaaclab_arena/relations/placement_asset.py

@@ -12,6 +12,7 @@
 
 from isaaclab_arena.assets.asset import Asset
 from isaaclab_arena.relations.collision_mode import CollisionMode
+from isaaclab_arena.relations.collision_object import CollisionComponent
 from isaaclab_arena.relations.relations import IsAnchor, Relation, RelationBase, UnaryRelation
 from isaaclab_arena.utils.bounding_box import quaternion_to_90_deg_z_quarters
 from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange
@@ -146,3 +147,18 @@ def get_collision_mesh(self) -> trimesh.Trimesh | None:
 
         Concrete (not abstract) so assets without a mesh simply keep the ``None`` default.
         """
+
+    def get_relation_bounding_box(self) -> AxisAlignedBoundingBox:
+        """Return the root-relative bounds a solver constrains against; subclasses may narrow the plain bounding box."""
+        return self.get_bounding_box()
+
+    def get_collision_components(self) -> list[CollisionComponent]:
+        """Return the asset's collision sub-volumes; a simple asset yields one identity-posed component."""
+        return [
+            CollisionComponent(
+                name=self.name,
+                local_pose=Pose.identity(),
+                bounding_box=self.get_bounding_box(),
+                mesh=self.get_collision_mesh(),
+            )
+        ]

isaaclab_arena/tests/test_bounding_box.py

@@ -193,6 +193,67 @@ def test_bounding_box_multi_env_overlaps():
     assert result[1].item() is False
 
 
+def test_union_encloses_all_boxes():
+    """Union spans the component-wise min/max of disjoint boxes."""
+    a = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 1.0))
+    b = AxisAlignedBoundingBox(min_point=(2.0, -1.0, 0.0), max_point=(3.0, 0.0, 2.0))
+    result = AxisAlignedBoundingBox.union([a, b])
+    torch.testing.assert_close(result.min_point, torch.tensor([[0.0, -1.0, 0.0]]))
+    torch.testing.assert_close(result.max_point, torch.tensor([[3.0, 1.0, 2.0]]))
+
+
+def test_union_broadcasts_single_box_over_batch():
+    """An N=1 box broadcasts against a batched box, yielding a per-env union."""
+    batched = AxisAlignedBoundingBox(
+        min_point=torch.tensor([[0.0, 0.0, 0.0], [5.0, 5.0, 0.0]]),
+        max_point=torch.tensor([[1.0, 1.0, 1.0], [6.0, 6.0, 1.0]]),
+    )
+    single = AxisAlignedBoundingBox(min_point=(-1.0, -1.0, -1.0), max_point=(0.0, 0.0, 0.0))
+    result = AxisAlignedBoundingBox.union([batched, single])
+    assert result.num_envs == 2
+    # The single box broadcasts into both rows, so check each independently.
+    torch.testing.assert_close(result.min_point[0], torch.tensor([-1.0, -1.0, -1.0]))
+    torch.testing.assert_close(result.max_point[0], torch.tensor([1.0, 1.0, 1.0]))
+    torch.testing.assert_close(result.min_point[1], torch.tensor([-1.0, -1.0, -1.0]))
+    torch.testing.assert_close(result.max_point[1], torch.tensor([6.0, 6.0, 1.0]))
+
+
+def test_union_rejects_mismatched_batch_sizes():
+    """Boxes with differing num_envs (neither N=1) are ambiguous and assert."""
+    two = 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]]),
+    )
+    three = AxisAlignedBoundingBox(
+        min_point=torch.zeros(3, 3),
+        max_point=torch.ones(3, 3),
+    )
+    with pytest.raises(AssertionError):
+        AxisAlignedBoundingBox.union([two, three])
+
+
+def test_union_rejects_empty_sequence():
+    """An empty sequence has no bounds to union and must assert, not raise a bare ValueError."""
+    with pytest.raises(AssertionError):
+        AxisAlignedBoundingBox.union([])
+
+
+def test_union_combines_two_per_env_boxes():
+    """Two genuinely per-env N=2 boxes union independently per row (no broadcasting)."""
+    a = AxisAlignedBoundingBox(
+        min_point=torch.tensor([[0.0, 0.0, 0.0], [5.0, 5.0, 5.0]]),
+        max_point=torch.tensor([[1.0, 1.0, 1.0], [6.0, 6.0, 6.0]]),
+    )
+    b = AxisAlignedBoundingBox(
+        min_point=torch.tensor([[-1.0, 0.0, 0.0], [5.0, 4.0, 5.0]]),
+        max_point=torch.tensor([[1.0, 2.0, 1.0], [7.0, 6.0, 6.0]]),
+    )
+    result = AxisAlignedBoundingBox.union([a, b])
+    assert result.num_envs == 2
+    torch.testing.assert_close(result.min_point, torch.tensor([[-1.0, 0.0, 0.0], [5.0, 4.0, 5.0]]))
+    torch.testing.assert_close(result.max_point, torch.tensor([[1.0, 2.0, 1.0], [7.0, 6.0, 6.0]]))
+
+
 def test_bounding_box_multi_env_get_corners_at():
     """Multi-env: get_corners_at() returns (N, 8, 3) tensor."""
     aabb = AxisAlignedBoundingBox(

isaaclab_arena/tests/test_embodiment_collision_mesh.py

@@ -3,6 +3,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import torch
 import traceback
 
 from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function
@@ -48,5 +49,49 @@ def test_embodiment_provides_robot_collision_mesh():
     assert result, f"Test {test_embodiment_provides_robot_collision_mesh.__name__} failed"
 
 
+def _test_droid_exposes_robot_and_stand_compound_geometry(simulation_app) -> bool:
+    """Check Droid splits its geometry into robot + stand components and uses the stand for relations."""
+
+    from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment
+    from isaaclab_arena.utils.pose import Pose
+
+    try:
+        emb = DroidAbsoluteJointPositionEmbodiment()
+
+        components = emb.get_collision_components()
+        assert [c.name for c in components] == ["robot", "stand"], f"unexpected components: {components}"
+        # Both prims are placed at the solved base pose, so each component sits at the Droid root.
+        assert all(c.local_pose == Pose.identity() for c in components), "components must share the root frame"
+
+        robot_component, stand_component = components
+        assert robot_component.mesh is not None, "robot component must carry its mesh for MESH-mode collision"
+
+        # The relation bbox is the stand footprint, not the arm's wider envelope: On/NextTo measure the
+        # base that sits by the support surface.
+        relation_bbox = emb.get_relation_bounding_box()
+        torch.testing.assert_close(relation_bbox.min_point, stand_component.bounding_box.min_point)
+        torch.testing.assert_close(relation_bbox.max_point, stand_component.bounding_box.max_point)
+
+        # The stand footprint must differ from the plain (arm) bounding box, otherwise the override is a no-op.
+        arm_bbox = emb.get_bounding_box()
+        relation_size = (relation_bbox.max_point - relation_bbox.min_point)[0]
+        arm_size = (arm_bbox.max_point - arm_bbox.min_point)[0]
+        assert not torch.allclose(relation_size, arm_size), "stand footprint should differ from the arm bbox"
+
+    except Exception as e:
+        print(f"Error: {e}")
+        traceback.print_exc()
+        return False
+
+    return True
+
+
+def test_droid_exposes_robot_and_stand_compound_geometry():
+    """Pytest entry point for the Droid compound-geometry test."""
+    result = run_simulation_app_function(_test_droid_exposes_robot_and_stand_compound_geometry, headless=True)
+    assert result, f"Test {test_droid_exposes_robot_and_stand_compound_geometry.__name__} failed"
+
+
 if __name__ == "__main__":
     test_embodiment_provides_robot_collision_mesh()
+    test_droid_exposes_robot_and_stand_compound_geometry()

isaaclab_arena/tests/test_relation_solver_embodiment.py

@@ -37,6 +37,21 @@ def _make_floor_and_robot():
     return floor, robot
 
 
+def test_placement_asset_defaults_to_single_component():
+    """A simple asset exposes one identity-posed component and a relation bbox equal to its bbox."""
+    bbox = AxisAlignedBoundingBox(min_point=(-0.2, -0.2, 0.0), max_point=(0.2, 0.2, 1.2))
+    robot = DummyEmbodiment(name="robot", bounding_box=bbox)
+
+    assert robot.get_relation_bounding_box() is bbox
+
+    components = robot.get_collision_components()
+    assert len(components) == 1
+    component = components[0]
+    assert component.local_pose == Pose.identity()
+    assert component.bounding_box is bbox
+    assert component.mesh is None
+
+
 def test_relation_solver_places_embodiment():
     floor, robot = _make_floor_and_robot()
 

isaaclab_arena/utils/bounding_box.py

@@ -13,6 +13,7 @@
 """
 
 import torch
+from collections.abc import Sequence
 
 from isaaclab_arena.utils.pose import Pose
 
@@ -45,6 +46,27 @@ def __getitem__(self, idx: int) -> "AxisAlignedBoundingBox":
             min_point=self._min_point[idx : idx + 1], max_point=self._max_point[idx : idx + 1]
         )
 
+    @staticmethod
+    def union(boxes: Sequence["AxisAlignedBoundingBox"]) -> "AxisAlignedBoundingBox":
+        """Return the tightest box enclosing every box in a common frame (component-wise min/max).
+
+        Args:
+            boxes: Non-empty sequence of boxes sharing num_envs, except N=1 boxes which broadcast.
+
+        Returns:
+            A box with num_envs = max over inputs, enclosing the union of all inputs.
+        """
+        assert len(boxes) > 0, "union requires at least one bounding box."
+        num_envs = max(box.num_envs for box in boxes)
+        for box in boxes:
+            assert box.num_envs in (
+                1,
+                num_envs,
+            ), f"union requires boxes with matching num_envs or N=1; got {box.num_envs} vs {num_envs}."
+        mins = torch.stack([box._min_point.expand(num_envs, 3) for box in boxes], dim=0)
+        maxs = torch.stack([box._max_point.expand(num_envs, 3) for box in boxes], dim=0)
+        return AxisAlignedBoundingBox(min_point=mins.amin(dim=0), max_point=maxs.amax(dim=0))
+
     @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."""