diff --git a/isaaclab_arena/assets/object.py b/isaaclab_arena/assets/object.py index 30322177..f55b7520 100644 --- a/isaaclab_arena/assets/object.py +++ b/isaaclab_arena/assets/object.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: Apache-2.0 +from typing import TYPE_CHECKING, Union + from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg @@ -11,10 +13,13 @@ from isaaclab_arena.utils.pose import Pose from isaaclab_arena.utils.usd_helpers import has_light, open_stage +if TYPE_CHECKING: + from isaaclab_arena.assets.relations import Relation + class Object(ObjectBase): """ - Encapsulates the pick-up object config for a pick-and-place environment. + A general-purpose object wrapper that encapsulates asset configurations for simulation environments. """ def __init__( @@ -24,7 +29,7 @@ def __init__( object_type: ObjectType | None = None, usd_path: str | None = None, scale: tuple[float, float, float] = (1.0, 1.0, 1.0), - initial_pose: Pose | None = None, + initial_pose: Union[Pose, "Relation", None] = None, **kwargs, ): if object_type is not ObjectType.SPAWNER: @@ -38,9 +43,12 @@ def __init__( self.initial_pose = initial_pose self.object_cfg = self._init_object_cfg() - def set_initial_pose(self, pose: Pose) -> None: + def set_initial_pose(self, pose: Union[Pose, "Relation"]) -> None: self.initial_pose = pose - self.object_cfg = self._add_initial_pose_to_cfg(self.object_cfg) + + # TODO(cvolk): How to do it properly? + # TODO(cvolk): Does the object_cfg need the initial pose here already? + # self.object_cfg = self._add_initial_pose_to_cfg(self.object_cfg) def get_initial_pose(self) -> Pose | None: return self.initial_pose diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py index 39ddbdaf..2ecc9f44 100644 --- a/isaaclab_arena/assets/object_base.py +++ b/isaaclab_arena/assets/object_base.py @@ -6,7 +6,7 @@ import torch from abc import ABC, abstractmethod from enum import Enum -from typing import Any +from typing import TYPE_CHECKING, Any from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg from isaaclab.envs import ManagerBasedEnv @@ -14,6 +14,9 @@ from isaaclab_arena.assets.asset import Asset +if TYPE_CHECKING: + from isaaclab_arena.utils.pose import Pose + class ObjectType(Enum): BASE = "base" @@ -110,3 +113,148 @@ def _generate_base_cfg(self) -> AssetBaseCfg: def _generate_spawner_cfg(self) -> AssetBaseCfg: # Object Subclasses must implement this method pass + + # Spatial Relationship Methods + def on_top_of( + self, + target: "ObjectBase", + clearance: float = 0.0, + x_offset: float = 0.0, + y_offset: float = 0.0, + ) -> "ObjectBase": + """ + Place this object on top of a target object. + + This method automatically computes the appropriate pose to place this object + on top of the target object, accounting for both objects' geometries. + + Args: + target: The target object to place this object on top of. + clearance: Additional vertical clearance between objects (default: 0.0). + x_offset: Horizontal offset in x direction from center (default: 0.0). + y_offset: Horizontal offset in y direction from center (default: 0.0). + + Returns: + Self, to allow method chaining. + + Example: + ```python + table = asset_registry.get_asset_by_name("table")() + box = asset_registry.get_asset_by_name("cracker_box")() + box.on_top_of(table) + scene = Scene(assets=[table, box]) + ``` + """ + from isaaclab_arena.utils.spatial_relationships import compute_bounding_box_from_usd, compute_on_top_of_pose + + # Get bounding boxes for both objects + # We need to access the usd_path and scale from the concrete Object class + # For now, we'll use getattr to handle this polymorphically + object_usd_path = getattr(self, "usd_path", None) + object_scale = getattr(self, "scale", (1.0, 1.0, 1.0)) + target_usd_path = getattr(target, "usd_path", None) + target_scale = getattr(target, "scale", (1.0, 1.0, 1.0)) + + if object_usd_path is None: + raise ValueError(f"Object {self.name} does not have a usd_path attribute") + if target_usd_path is None: + raise ValueError(f"Target object {target.name} does not have a usd_path attribute") + + # Get the target's current pose (if set) + target_pose = getattr(target, "initial_pose", None) + + # Compute bounding boxes + object_bbox = compute_bounding_box_from_usd(object_usd_path, scale=object_scale) + target_bbox = compute_bounding_box_from_usd(target_usd_path, scale=target_scale, pose=target_pose) + + # Compute the placement pose + placement_pose = compute_on_top_of_pose( + object_bbox=object_bbox, + target_bbox=target_bbox, + clearance=clearance, + x_offset=x_offset, + y_offset=y_offset, + ) + + # Set the initial pose on this object + self.set_initial_pose(placement_pose) + + def next_to( + self, + target: "ObjectBase", + side: str = "right", + clearance: float = 0.01, + align_bottom: bool = True, + ) -> "ObjectBase": + """ + Place this object next to a target object. + + This method automatically computes the appropriate pose to place this object + beside the target object, accounting for both objects' geometries. + + **Important**: Directions are in the world coordinate frame, not relative to + the target object's orientation: + - "right" = -Y world direction + - "left" = +Y world direction + - "front" = -X world direction + - "back" = +X world direction + + Args: + target: The target object to place this object next to. + side: Which side to place the object ("left", "right", "front", "back"). + These directions are in world frame, not relative to target's orientation. + clearance: Horizontal clearance between objects (default: 0.01). + align_bottom: If True, align bottoms; if False, center vertically (default: True). + + Returns: + Self, to allow method chaining. + + Example: + ```python + laptop = asset_registry.get_asset_by_name("laptop")() + mug = asset_registry.get_asset_by_name("mug")() + # Places mug in -Y direction from laptop (world frame) + mug.next_to(laptop, side="right") + scene = Scene(assets=[laptop, mug]) + ``` + + Note: + This is a limitation of the MVP. Future versions may support + placement relative to the target object's local coordinate frame. + """ + from isaaclab_arena.utils.spatial_relationships import compute_bounding_box_from_usd, compute_next_to_pose + + # Get bounding boxes for both objects + object_usd_path = getattr(self, "usd_path", None) + object_scale = getattr(self, "scale", (1.0, 1.0, 1.0)) + target_usd_path = getattr(target, "usd_path", None) + target_scale = getattr(target, "scale", (1.0, 1.0, 1.0)) + + if object_usd_path is None: + raise ValueError(f"Object {self.name} does not have a usd_path attribute") + if target_usd_path is None: + raise ValueError(f"Target object {target.name} does not have a usd_path attribute") + + # Get the target's current pose (if set) + target_pose = getattr(target, "initial_pose", None) + + # Compute bounding boxes + object_bbox = compute_bounding_box_from_usd(object_usd_path, scale=object_scale) + target_bbox = compute_bounding_box_from_usd(target_usd_path, scale=target_scale, pose=target_pose) + + # Compute the placement pose + placement_pose = compute_next_to_pose( + object_bbox=object_bbox, + target_bbox=target_bbox, + side=side, + clearance=clearance, + align_bottom=align_bottom, + ) + + # Set the initial pose on this object + self.set_initial_pose(placement_pose) + + @abstractmethod + def set_initial_pose(self, pose: "Pose") -> None: + """Set the initial pose of the object. Must be implemented by subclasses.""" + pass diff --git a/isaaclab_arena/assets/relations.py b/isaaclab_arena/assets/relations.py new file mode 100644 index 00000000..21d000ac --- /dev/null +++ b/isaaclab_arena/assets/relations.py @@ -0,0 +1,51 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from abc import abstractmethod +from typing import TYPE_CHECKING + +from isaaclab_arena.assets.asset import Asset +from isaaclab_arena.utils.pose import Pose + +# Use TYPE_CHECKING to avoid circular import at runtime +if TYPE_CHECKING: + from isaaclab_arena.scene.scene import Scene + + +class Relation: + def __init__(self, parent_asset: Asset, child_asset: Asset): + self.parent_asset: Asset = parent_asset + self.child_asset: Asset = child_asset + + @abstractmethod + def resolve(self): + pass + + +class OnRelation(Relation): + def resolve(self): + # Get the pose of the parent + # Resolve the pose of the child relative to the parent + # Add to world frame + # return the pose in world frame + print(f"Resolving on relation between {self.parent_asset.name} and {self.child_asset.name}") + return Pose(position_xyz=(0.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + + +class RelationResolver: + def __init__(self, scene: "Scene"): + self.scene = scene + + # This would actually set the initial_poses + def resolve_relations(self): + # Get assets + for asset in self.scene.assets.values(): + if isinstance(asset.initial_pose, Relation): + print(f"Asset {asset.name} has initial pose of type {type(asset.initial_pose)}") + pose = asset.initial_pose.resolve() + print(f"Now setting initial pose of asset {asset.name} to {pose}") + asset.set_initial_pose(pose) + else: + print(f"Asset {asset.name} has no initial pose") diff --git a/isaaclab_arena/examples/compile_env_notebook.py b/isaaclab_arena/examples/compile_env_notebook.py index 9ebe8611..61e5201a 100644 --- a/isaaclab_arena/examples/compile_env_notebook.py +++ b/isaaclab_arena/examples/compile_env_notebook.py @@ -15,6 +15,7 @@ simulation_app = AppLauncher() from isaaclab_arena.assets.asset_registry import AssetRegistry +from isaaclab_arena.assets.relations import OnRelation, RelationResolver from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment @@ -27,10 +28,15 @@ background = asset_registry.get_asset_by_name("kitchen")() embodiment = asset_registry.get_asset_by_name("franka")() cracker_box = asset_registry.get_asset_by_name("cracker_box")() - +tomato_soup_can = asset_registry.get_asset_by_name("tomato_soup_can")() cracker_box.set_initial_pose(Pose(position_xyz=(0.4, 0.0, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) -scene = Scene(assets=[background, cracker_box]) +print("Setting initial pose of tomato soup can to on relation with cracker box") +tomato_soup_can.set_initial_pose(OnRelation(tomato_soup_can, cracker_box)) + +scene = Scene(assets=[background, cracker_box, tomato_soup_can]) + + isaaclab_arena_environment = IsaacLabArenaEnvironment( name="reference_object_test", embodiment=embodiment, @@ -54,3 +60,6 @@ env.step(actions) # %% + +# TODO(cvolk) +# We still need an anchor point diff --git a/isaaclab_arena/examples/compile_env_notebook_v2.py b/isaaclab_arena/examples/compile_env_notebook_v2.py new file mode 100644 index 00000000..017b6598 --- /dev/null +++ b/isaaclab_arena/examples/compile_env_notebook_v2.py @@ -0,0 +1,66 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# %% + +import torch +import tqdm + +import pinocchio # noqa: F401 +from isaaclab.app import AppLauncher + +print("Launching simulation app once in notebook") +simulation_app = AppLauncher() + +from isaaclab_arena.assets.asset_registry import AssetRegistry +from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser +from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder +from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment +from isaaclab_arena.scene.scene import Scene +from isaaclab_arena.tasks.dummy_task import DummyTask +from isaaclab_arena.utils.pose import Pose + +asset_registry = AssetRegistry() + +background = asset_registry.get_asset_by_name("kitchen")() +embodiment = asset_registry.get_asset_by_name("franka")() +cracker_box = asset_registry.get_asset_by_name("cracker_box")() +cracker_box.set_initial_pose(Pose(position_xyz=(0.4, 0.0, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) +tomato_soup_can = asset_registry.get_asset_by_name("tomato_soup_can")() + +microwave = asset_registry.get_asset_by_name("microwave")() + +tomato_soup_can.next_to(cracker_box, side="right", clearance=0.05) +microwave.next_to(tomato_soup_can, side="right", clearance=0.05) +mustard_bottle.on_top_of(microwave) + +scene = Scene(assets=[background, cracker_box, tomato_soup_can, microwave, mustard_bottle]) + +isaaclab_arena_environment = IsaacLabArenaEnvironment( + name="reference_object_test", + embodiment=embodiment, + scene=scene, + task=DummyTask(), + teleop_device=None, +) + +args_cli = get_isaaclab_arena_cli_parser().parse_args([]) +env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) +env = env_builder.make_registered() +env.reset() + +# %% + +# Run some zero actions. +NUM_STEPS = 1000 +for _ in tqdm.tqdm(range(NUM_STEPS)): + with torch.inference_mode(): + actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) + env.step(actions) + +# %% + +# TODO(cvolk) +# We still need an anchor point diff --git a/isaaclab_arena/tests/test_spatial_relationships.py b/isaaclab_arena/tests/test_spatial_relationships.py new file mode 100644 index 00000000..fe43105d --- /dev/null +++ b/isaaclab_arena/tests/test_spatial_relationships.py @@ -0,0 +1,156 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Unit tests for the spatial relationships API. + +These tests verify that the spatial relationship functions correctly compute +bounding boxes and placement poses. +""" + +import pytest + +from isaaclab_arena.utils.spatial_relationships import BoundingBox, compute_next_to_pose, compute_on_top_of_pose + + +class TestBoundingBox: + """Test the BoundingBox dataclass.""" + + def test_size(self): + """Test size calculation.""" + bbox = BoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 2.0, 3.0)) + assert bbox.size == (1.0, 2.0, 3.0) + + def test_center(self): + """Test center calculation.""" + bbox = BoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 4.0, 6.0)) + assert bbox.center == (1.0, 2.0, 3.0) + + def test_top_surface_z(self): + """Test top surface z coordinate.""" + bbox = BoundingBox(min_point=(0.0, 0.0, 1.0), max_point=(1.0, 1.0, 3.0)) + assert bbox.top_surface_z == 3.0 + + def test_bottom_surface_z(self): + """Test bottom surface z coordinate.""" + bbox = BoundingBox(min_point=(0.0, 0.0, 1.0), max_point=(1.0, 1.0, 3.0)) + assert bbox.bottom_surface_z == 1.0 + + +class TestComputeOnTopOfPose: + """Test the compute_on_top_of_pose function.""" + + def test_basic_placement(self): + """Test basic on-top placement.""" + # Create a 1x1x1 cube as the object + object_bbox = BoundingBox(min_point=(-0.5, -0.5, -0.5), max_point=(0.5, 0.5, 0.5)) + # Create a 2x2x1 surface as the target (centered at origin) + target_bbox = BoundingBox(min_point=(-1.0, -1.0, -0.5), max_point=(1.0, 1.0, 0.5)) + + pose = compute_on_top_of_pose(object_bbox, target_bbox) + + # Object should be centered on target horizontally + assert pose.position_xyz[0] == 0.0 + assert pose.position_xyz[1] == 0.0 + # Object bottom (at -0.5 relative to center) should touch target top (at 0.5) + # So object center should be at 0.5 + 0.5 = 1.0 + assert pose.position_xyz[2] == 1.0 + # No rotation + assert pose.rotation_wxyz == (1.0, 0.0, 0.0, 0.0) + + def test_placement_with_clearance(self): + """Test placement with additional clearance.""" + object_bbox = BoundingBox(min_point=(-0.5, -0.5, -0.5), max_point=(0.5, 0.5, 0.5)) + target_bbox = BoundingBox(min_point=(-1.0, -1.0, -0.5), max_point=(1.0, 1.0, 0.5)) + + pose = compute_on_top_of_pose(object_bbox, target_bbox, clearance=0.1) + + # Z position should include the clearance + assert pose.position_xyz[2] == 1.1 + + def test_placement_with_offsets(self): + """Test placement with x/y offsets.""" + object_bbox = BoundingBox(min_point=(-0.5, -0.5, -0.5), max_point=(0.5, 0.5, 0.5)) + target_bbox = BoundingBox(min_point=(-1.0, -1.0, -0.5), max_point=(1.0, 1.0, 0.5)) + + pose = compute_on_top_of_pose(object_bbox, target_bbox, x_offset=0.3, y_offset=-0.2) + + assert pose.position_xyz[0] == 0.3 + assert pose.position_xyz[1] == -0.2 + + +class TestComputeNextToPose: + """Test the compute_next_to_pose function.""" + + def test_placement_right(self): + """Test placement to the right.""" + object_bbox = BoundingBox(min_point=(-0.5, -0.5, -0.5), max_point=(0.5, 0.5, 0.5)) + target_bbox = BoundingBox(min_point=(-1.0, -1.0, -0.5), max_point=(1.0, 1.0, 0.5)) + + pose = compute_next_to_pose(object_bbox, target_bbox, side="right", clearance=0.0) + + # right = -Y direction + # Target bottom edge (in Y) is at -1.0, object top edge should touch it + # Object center should be at -1.0 - 0.5 = -1.5 + assert pose.position_xyz[0] == 0.0 # Centered in X + assert pose.position_xyz[1] == -1.5 # At -Y + assert pose.position_xyz[2] == 0.0 # Bottom aligned + + def test_placement_left(self): + """Test placement to the left.""" + object_bbox = BoundingBox(min_point=(-0.5, -0.5, -0.5), max_point=(0.5, 0.5, 0.5)) + target_bbox = BoundingBox(min_point=(-1.0, -1.0, -0.5), max_point=(1.0, 1.0, 0.5)) + + pose = compute_next_to_pose(object_bbox, target_bbox, side="left", clearance=0.0) + + # left = +Y direction + # Target top edge (in Y) is at 1.0, object bottom edge should touch it + # Object center should be at 1.0 + 0.5 = 1.5 + assert pose.position_xyz[0] == 0.0 # Centered in X + assert pose.position_xyz[1] == 1.5 # At +Y + + def test_placement_front(self): + """Test placement in front.""" + object_bbox = BoundingBox(min_point=(-0.5, -0.5, -0.5), max_point=(0.5, 0.5, 0.5)) + target_bbox = BoundingBox(min_point=(-1.0, -1.0, -0.5), max_point=(1.0, 1.0, 0.5)) + + pose = compute_next_to_pose(object_bbox, target_bbox, side="front", clearance=0.0) + + # front = -X direction + # Target left edge (in X) is at -1.0, object right edge should touch it + # Object center should be at -1.0 - 0.5 = -1.5 + assert pose.position_xyz[0] == -1.5 # At -X + assert pose.position_xyz[1] == 0.0 # Centered in Y + + def test_placement_back(self): + """Test placement behind.""" + object_bbox = BoundingBox(min_point=(-0.5, -0.5, -0.5), max_point=(0.5, 0.5, 0.5)) + target_bbox = BoundingBox(min_point=(-1.0, -1.0, -0.5), max_point=(1.0, 1.0, 0.5)) + + pose = compute_next_to_pose(object_bbox, target_bbox, side="back", clearance=0.0) + + # back = +X direction + # Target right edge (in X) is at 1.0, object left edge should touch it + # Object center should be at 1.0 + 0.5 = 1.5 + assert pose.position_xyz[0] == 1.5 # At +X + assert pose.position_xyz[1] == 0.0 # Centered in Y + + def test_placement_with_clearance(self): + """Test placement with clearance.""" + object_bbox = BoundingBox(min_point=(-0.5, -0.5, -0.5), max_point=(0.5, 0.5, 0.5)) + target_bbox = BoundingBox(min_point=(-1.0, -1.0, -0.5), max_point=(1.0, 1.0, 0.5)) + + pose = compute_next_to_pose(object_bbox, target_bbox, side="right", clearance=0.1) + + # right = -Y, with 0.1 clearance, object center should be at -1.0 - 0.1 - 0.5 = -1.6 + assert pose.position_xyz[1] == -1.6 + + def test_invalid_side(self): + """Test that invalid side raises ValueError.""" + object_bbox = BoundingBox(min_point=(-0.5, -0.5, -0.5), max_point=(0.5, 0.5, 0.5)) + target_bbox = BoundingBox(min_point=(-1.0, -1.0, -0.5), max_point=(1.0, 1.0, 0.5)) + + with pytest.raises(ValueError, match="Invalid side"): + compute_next_to_pose(object_bbox, target_bbox, side="invalid") diff --git a/isaaclab_arena/utils/spatial_relationships.py b/isaaclab_arena/utils/spatial_relationships.py new file mode 100644 index 00000000..964e58e8 --- /dev/null +++ b/isaaclab_arena/utils/spatial_relationships.py @@ -0,0 +1,237 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# 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 +""" + +from dataclasses import dataclass + +from pxr import Gf, Usd, UsdGeom + +from isaaclab_arena.utils.pose import Pose + + +@dataclass +class BoundingBox: + """Represents an axis-aligned bounding box in 3D space.""" + + min_point: tuple[float, float, float] + """Minimum point (x, y, z) of the bounding box.""" + + max_point: tuple[float, float, float] + """Maximum point (x, y, z) of the bounding box.""" + + @property + def size(self) -> tuple[float, float, float]: + """Returns the size (width, depth, height) of the bounding box.""" + return ( + self.max_point[0] - self.min_point[0], + self.max_point[1] - self.min_point[1], + self.max_point[2] - self.min_point[2], + ) + + @property + def center(self) -> tuple[float, float, float]: + """Returns the center point of the bounding box.""" + return ( + (self.min_point[0] + self.max_point[0]) / 2.0, + (self.min_point[1] + self.max_point[1]) / 2.0, + (self.min_point[2] + self.max_point[2]) / 2.0, + ) + + @property + def top_surface_z(self) -> float: + """Returns the z-coordinate of the top surface.""" + return self.max_point[2] + + @property + def bottom_surface_z(self) -> float: + """Returns the z-coordinate of the bottom surface.""" + return self.min_point[2] + + +def compute_bounding_box_from_usd( + usd_path: str, + scale: tuple[float, float, float] = (1.0, 1.0, 1.0), + pose: Pose | None = None, +) -> BoundingBox: + """ + Compute the world-space bounding box of a USD asset. + + Args: + usd_path: Path to the USD file. + scale: Scale to apply to the asset (x, y, z). + pose: Optional pose of the asset. If None, uses identity pose. + + Returns: + BoundingBox containing the min and max points in world space. + """ + # Open the USD stage + stage = Usd.Stage.Open(usd_path) + if not stage: + raise ValueError(f"Failed to open USD file: {usd_path}") + + # Get the default prim (or pseudo root if no default prim) + default_prim = stage.GetDefaultPrim() + if not default_prim: + default_prim = stage.GetPseudoRoot() + + # Compute the bounding box using USD's built-in functionality + # This computes the bounding box in the local space of the prim + bbox_cache = UsdGeom.BBoxCache(Usd.TimeCode.Default(), includedPurposes=[UsdGeom.Tokens.default_]) + bbox = bbox_cache.ComputeWorldBound(default_prim) + + # Get the range (bounding box) + bbox_range = bbox.ComputeAlignedBox() + min_point = bbox_range.GetMin() + max_point = bbox_range.GetMax() + + # Apply scale + min_point = Gf.Vec3d(min_point[0] * scale[0], min_point[1] * scale[1], min_point[2] * scale[2]) + max_point = Gf.Vec3d(max_point[0] * scale[0], max_point[1] * scale[1], max_point[2] * scale[2]) + + # Apply pose transformation if provided + if pose is not None: + # Transform the bounding box corners by the pose + # For simplicity in MVP, we'll just translate the bbox by the position + # A more sophisticated implementation would rotate the bbox as well + min_point = Gf.Vec3d( + min_point[0] + pose.position_xyz[0], + min_point[1] + pose.position_xyz[1], + min_point[2] + pose.position_xyz[2], + ) + max_point = Gf.Vec3d( + max_point[0] + pose.position_xyz[0], + max_point[1] + pose.position_xyz[1], + max_point[2] + pose.position_xyz[2], + ) + + return BoundingBox( + min_point=(min_point[0], min_point[1], min_point[2]), + max_point=(max_point[0], max_point[1], max_point[2]), + ) + + +def compute_on_top_of_pose( + object_bbox: BoundingBox, + target_bbox: BoundingBox, + clearance: float = 0.0, + x_offset: float = 0.0, + y_offset: float = 0.0, +) -> Pose: + """ + Compute a pose that places an object on top of a target object. + + The object will be placed such that its bottom surface sits on the target's top surface, + centered on the target (unless offsets are provided). + + Args: + object_bbox: Bounding box of the object to be placed. + target_bbox: Bounding box of the target (base) object. + clearance: Additional vertical clearance between objects (default: 0.0). + x_offset: Horizontal offset in x direction from center (default: 0.0). + y_offset: Horizontal offset in y direction from center (default: 0.0). + + Returns: + Pose for the object that places it on top of the target. + """ + # Calculate the z position: target's top + clearance + object's height/2 + # We need to account for the object's center being at its geometric center + object_half_height = object_bbox.size[2] / 2.0 + + # The bottom of the object should be at the top of the target + z_position = target_bbox.top_surface_z + clearance + object_half_height + + # Center the object on the target (with optional offsets) + x_position = target_bbox.center[0] + x_offset + y_position = target_bbox.center[1] + y_offset + + # Return pose with identity rotation (no rotation by default) + return Pose( + position_xyz=(x_position, y_position, z_position), + rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + ) + + +def compute_next_to_pose( + object_bbox: BoundingBox, + target_bbox: BoundingBox, + side: str = "right", + clearance: float = 0.01, + align_bottom: bool = True, +) -> Pose: + """ + Compute a pose that places an object next to a target object. + + **Important**: Directions are defined in the **world coordinate frame**, not relative + to the target object's orientation: + - "right" = -Y direction in world frame + - "left" = +Y direction in world frame + - "front" = -X direction in world frame + - "back" = +X direction in world frame + + This means the placement does NOT account for the target object's rotation. + This is a limitation of the MVP implementation. + + Args: + object_bbox: Bounding box of the object to be placed. + target_bbox: Bounding box of the target object. + side: Which side to place the object ("left", "right", "front", "back"). + Directions are in world frame, not relative to target's orientation. + clearance: Horizontal clearance between objects (default: 0.01). + align_bottom: If True, align bottoms; if False, center vertically (default: True). + + Returns: + Pose for the object that places it next to the target. + + Note: + For rotated objects, "right" will still be -Y in world coordinates, + which may not correspond to the intuitive "right side" of the object. + """ + # Calculate the base z position + if align_bottom: + # Align the bottom surfaces + object_half_height = object_bbox.size[2] / 2.0 + z_position = target_bbox.bottom_surface_z + object_half_height + else: + # Center vertically with the target + z_position = target_bbox.center[2] + + # Calculate horizontal position based on side + # Convention: right = -Y, left = +Y, front = -X, back = +X + object_half_width = object_bbox.size[0] / 2.0 + object_half_depth = object_bbox.size[1] / 2.0 + target_half_width = target_bbox.size[0] / 2.0 + target_half_depth = target_bbox.size[1] / 2.0 + + if side == "right": + # right = -Y + x_position = target_bbox.center[0] + y_position = target_bbox.center[1] - target_half_depth - clearance - object_half_depth + elif side == "left": + # left = +Y + x_position = target_bbox.center[0] + y_position = target_bbox.center[1] + target_half_depth + clearance + object_half_depth + elif side == "front": + # front = -X + x_position = target_bbox.center[0] - target_half_width - clearance - object_half_width + y_position = target_bbox.center[1] + elif side == "back": + # back = +X + x_position = target_bbox.center[0] + target_half_width + clearance + object_half_width + y_position = target_bbox.center[1] + else: + raise ValueError(f"Invalid side: {side}. Must be 'left', 'right', 'front', or 'back'.") + + return Pose( + position_xyz=(x_position, y_position, z_position), + rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + )