diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index ee3686d03c..69bc611bf4 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/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 diff --git a/isaaclab_arena/relations/collision_object.py b/isaaclab_arena/relations/collision_object.py index 115670f0a1..87ad5e999c 100644 --- a/isaaclab_arena/relations/collision_object.py +++ b/isaaclab_arena/relations/collision_object.py @@ -8,6 +8,7 @@ from __future__ import annotations import trimesh +from dataclasses import dataclass from typing import Protocol from isaaclab_arena.relations.collision_mode import CollisionMode @@ -15,6 +16,30 @@ 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.""" diff --git a/isaaclab_arena/relations/placement_asset.py b/isaaclab_arena/relations/placement_asset.py index fd31535470..94e5ae5fbe 100644 --- a/isaaclab_arena/relations/placement_asset.py +++ b/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(), + ) + ] diff --git a/isaaclab_arena/tests/test_bounding_box.py b/isaaclab_arena/tests/test_bounding_box.py index 7d38959e3e..c949b35857 100644 --- a/isaaclab_arena/tests/test_bounding_box.py +++ b/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( diff --git a/isaaclab_arena/tests/test_embodiment_collision_mesh.py b/isaaclab_arena/tests/test_embodiment_collision_mesh.py index 9b9487d2e6..4bbf072a4c 100644 --- a/isaaclab_arena/tests/test_embodiment_collision_mesh.py +++ b/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() diff --git a/isaaclab_arena/tests/test_relation_solver_embodiment.py b/isaaclab_arena/tests/test_relation_solver_embodiment.py index 9fc2484097..8419281883 100644 --- a/isaaclab_arena/tests/test_relation_solver_embodiment.py +++ b/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() diff --git a/isaaclab_arena/utils/bounding_box.py b/isaaclab_arena/utils/bounding_box.py index bac96332fa..0559d81927 100644 --- a/isaaclab_arena/utils/bounding_box.py +++ b/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."""