Integrate Droid lightwheel kitchen example with robot placement

isaaclab_arena/assets/asset.py

@@ -47,3 +47,7 @@ def get_variation(self, name: str) -> VariationBase:
     def get_variations(self) -> list[VariationBase]:
         """Return every variation attached to this asset, enabled or not."""
         return list(self.variations.values())
+
+    def get_scene_name(self) -> str:
+        """Return the Isaac Lab scene key for the asset."""
+        return self.name

isaaclab_arena/assets/object.py

@@ -5,10 +5,7 @@
 from __future__ import annotations
 
 import torch
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    import trimesh
+from typing import Any
 
 from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg
 from isaaclab.sensors.contact_sensor.contact_sensor_cfg import ContactSensorCfg
@@ -18,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, quaternion_to_90_deg_z_quarters
+from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 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
@@ -66,33 +63,13 @@ def __init__(
         self.object_cfg = self._init_object_cfg()
         self.event_cfg = self._init_event_cfg()
 
-    def add_relation(self, relation: RelationBase) -> None:
-        """Add a relation to this object."""
-        self.relations.append(relation)
-
     def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get local bounding box (relative to object origin)."""
         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
 
-    def get_collision_mesh(self) -> trimesh.Trimesh | None:
-        """Return None: USD-backed objects expose no preloaded collision mesh."""
-
-    def get_world_bounding_box(self) -> AxisAlignedBoundingBox:
-        """Get bounding box in world coordinates (local bbox rotated and translated).
-
-        Only 90° rotations around Z axis are supported. An assertion error is raised
-        for any other rotation. If initial_pose is a PoseRange (not a fixed Pose),
-        returns the local bounding box without transformation.
-        """
-        local_bbox = self.get_bounding_box()
-        if self.initial_pose is None or not isinstance(self.initial_pose, Pose):
-            return local_bbox
-        quarters = quaternion_to_90_deg_z_quarters(self.initial_pose.rotation_xyzw)
-        return local_bbox.rotated_90_around_z(quarters).translated(self.initial_pose.position_xyz)
-
     def get_corners(self, pos: torch.Tensor) -> torch.Tensor:
         assert self.usd_path is not None
         if self.bounding_box is None:

isaaclab_arena/assets/object_base.py

@@ -7,29 +7,21 @@
 
 import torch
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
 
 import warp as wp
 from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg
-
-if TYPE_CHECKING:
-    import trimesh
 from isaaclab.envs import ManagerBasedEnv
 from isaaclab.managers import EventTermCfg, SceneEntityCfg
 from isaaclab.sensors.contact_sensor.contact_sensor_cfg import ContactSensorCfg
 from isaaclab_tasks.manager_based.manipulation.stack.mdp.franka_stack_events import randomize_object_pose
 
-from isaaclab_arena.assets.asset import Asset
-
 # Re-export ObjectType from the lightweight module so existing
 # `from isaaclab_arena.assets.object_base import ObjectType` consumers keep working,
 # while pure-Python spec modules can import from `object_type` directly without
 # pulling in isaaclab/omni/pxr at module-load time.
 from isaaclab_arena.assets.object_type import ObjectType
-from isaaclab_arena.relations.collision_mode import CollisionMode
-from isaaclab_arena.relations.relations import IsAnchor, Relation, RelationBase, UnaryRelation
+from isaaclab_arena.relations.placement_asset import PlacementAsset
 from isaaclab_arena.terms.events import set_object_pose, set_object_pose_per_env
-from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange
 from isaaclab_arena.utils.velocity import Velocity
 from isaaclab_arena.variations.object_mass_variation import ObjectMassVariation
@@ -40,7 +32,7 @@
 ]
 
 
-class ObjectBase(Asset, ABC):
+class ObjectBase(PlacementAsset, ABC):
     """Parent class for (spawnable) Object and ObjectReference."""
 
     def __init__(
@@ -57,36 +49,9 @@ def __init__(
         self.object_type = object_type
         if self.object_type == ObjectType.RIGID:
             self.add_variation(ObjectMassVariation(self.name))
-        self.initial_pose: Pose | PoseRange | PosePerEnv | None = None
         self.initial_velocity: Velocity | None = None
         self.object_cfg = None
         self.event_cfg = None
-        self.relations: list[RelationBase] = []
-        # None means use the solver's default collision mode for this object.
-        self.collision_mode: CollisionMode | None = None
-        # If True, mesh collision replaces non-watertight meshes with their convex hull.
-        self.repair_collision_mesh_non_watertight = True
-
-    def get_initial_pose(self) -> Pose | PoseRange | PosePerEnv | None:
-        """Return the current initial pose of this object.
-
-        Subclasses may override to derive the pose from other sources
-        (e.g. a parent asset), falling back to ``self.initial_pose``.
-        """
-        return self.initial_pose
-
-    @abstractmethod
-    def get_bounding_box(self) -> AxisAlignedBoundingBox:
-        """Get local bounding box (relative to object origin)."""
-        ...
-
-    @abstractmethod
-    def get_world_bounding_box(self) -> AxisAlignedBoundingBox:
-        """Get bounding box in world coordinates (local bbox rotated and translated)."""
-        ...
-
-    def get_collision_mesh(self) -> trimesh.Trimesh | None:
-        """Return collision mesh, or None to fall back to AABB overlap."""
 
     def _get_initial_pose_as_pose(self) -> Pose | None:
         """Return a single ``Pose`` suitable for *init_state* and bounding-box calculations.
@@ -111,12 +76,25 @@ def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None:
             pose: A fixed ``Pose``, a ``PoseRange`` (randomised on reset),
                 or a ``PosePerEnv`` (distinct pose per environment).
         """
+        self._set_pose_state(pose)
+        self.event_cfg = self._init_event_cfg()
+
+    def _set_pose_state(self, pose: Pose | PoseRange | PosePerEnv) -> None:
+        """Update the stored pose and materialized object configuration."""
         self.initial_pose = pose
         initial_pose = self._get_initial_pose_as_pose()
         if initial_pose is not None and self.object_cfg is not None:
             self.object_cfg.init_state.pos = initial_pose.position_xyz
             self.object_cfg.init_state.rot = initial_pose.rotation_xyzw
-        self.event_cfg = self._init_event_cfg()
+
+    def set_spawn_pose(self, pose: Pose) -> None:
+        """Set the scene-construction pose without rebuilding the reset event."""
+        assert self.object_cfg is not None, "object_cfg must be initialized before setting the spawn pose"
+        self._set_pose_state(pose)
+
+    def has_pose_reset_event(self) -> bool:
+        """Return whether the asset owns a root-pose reset event."""
+        return self.event_cfg is not None
 
     def set_initial_velocity(self, velocity: Velocity) -> None:
         """Set / override the initial velocity and rebuild derived configs.
@@ -181,19 +159,6 @@ def _init_event_cfg(self) -> EventTermCfg | None:
                 },
             )
 
-    def get_relations(self) -> list[RelationBase]:
-        """Get all relations for this object."""
-        return self.relations
-
-    @property
-    def is_anchor(self) -> bool:
-        """True if this object has an IsAnchor relation."""
-        return any(isinstance(r, IsAnchor) for r in self.relations)
-
-    def get_spatial_relations(self) -> list[RelationBase]:
-        """Get only spatial relations (On, NextTo, AtPosition, etc.), excluding markers like IsAnchor."""
-        return [r for r in self.relations if isinstance(r, (Relation, UnaryRelation))]
-
     def set_prim_path(self, prim_path: str) -> None:
         self.prim_path = prim_path
 

isaaclab_arena/embodiments/droid/droid.py

@@ -40,11 +40,14 @@
 from isaaclab_arena.embodiments.droid.observations import arm_joint_pos, ee_pos, ee_quat, gripper_pos
 from isaaclab_arena.embodiments.embodiment_base import EmbodimentBase
 from isaaclab_arena.embodiments.franka.franka import franka_stack_events
+from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 from isaaclab_arena.utils.cameras import ArenaCameraCfg
-from isaaclab_arena.utils.pose import Pose, translate_by_xyz_offset
+from isaaclab_arena.utils.pose import Pose, PosePerEnv, compose_poses, translate_by_xyz_offset
 
 # The base stand's x/y footprint.
 _STAND_FOOTPRINT_SCALE_XY: tuple[float, float] = (1.2, 1.2)
+# Stand root offset in the robot base frame (matches ``DroidSceneCfg.stand`` default init_state).
+_STAND_ROOT_OFFSET_IN_ROBOT_FRAME: tuple[float, float, float] = (-0.05, 0.0, 0.0)
 # The default stand height.
 _DEFAULT_STAND_HEIGHT_M: float = 1.35
 _FALLBACK_STAND_UNIT_HEIGHT_M: float = 0.795
@@ -98,16 +101,55 @@ def __init__(
         self.mimic_env = None
         self.add_camera_variations(self.camera_config)
 
-    def set_initial_pose(self, pose: Pose) -> None:
-        """Store the requested base pose, lifted by the stand-height offset to match the spawned base."""
-        super().set_initial_pose(pose.translate(self._robot_base_offset))
+    def set_initial_pose(self, pose: Pose | PosePerEnv) -> None:
+        """Store the requested base pose(s), lifted by the stand-height offset to match the spawned base."""
+        if isinstance(pose, PosePerEnv):
+            super().set_initial_pose(PosePerEnv(poses=[p.translate(self._robot_base_offset) for p in pose.poses]))
+        else:
+            super().set_initial_pose(pose.translate(self._robot_base_offset))
+
+    def set_spawn_pose(self, pose: Pose) -> None:
+        """Set the scene-construction base pose, lifted by the stand-height offset like ``set_initial_pose``."""
+        super().set_spawn_pose(pose.translate(self._robot_base_offset))
+
+    def get_bounding_box(self) -> AxisAlignedBoundingBox:
+        """Return stand bounding box as proxy for the robot bounding box.
+
+        TODO(qianl): Hack to place the robot close to surfaces using only the stand bbox.
+        Remove after switching the embodiment to mesh collision mode.
+        """
+        from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd
+
+        assert self.scene_config is not None, "scene_config must be populated before placement"
+        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))
+        stand_bbox = compute_local_bounding_box_from_usd(spawn.usd_path, scale)
+        stand_offset = translate_by_xyz_offset(_STAND_ROOT_OFFSET_IN_ROBOT_FRAME, self._robot_base_offset)
+        return stand_bbox.translated(stand_offset)
+
+    def _stand_pose_from_robot_pose(self, robot_pose: Pose) -> Pose:
+        """Return the stand root pose for a lifted robot base pose."""
+        return compose_poses(
+            robot_pose,
+            Pose(position_xyz=_STAND_ROOT_OFFSET_IN_ROBOT_FRAME, rotation_xyzw=(0.0, 0.0, 0.0, 1.0)),
+        )
+
+    def layout_pose_to_scene_writes(self, layout_pose: Pose) -> list[tuple[str, Pose]]:
+        """Write the robot and stand roots, lifting the solver layout pose for sim spawn."""
+        robot_pose = layout_pose.translate(self._robot_base_offset)
+        stand_pose = self._stand_pose_from_robot_pose(robot_pose)
+        return [
+            (self.get_scene_name(), robot_pose),
+            ("stand", stand_pose),
+        ]
 
     def _update_scene_cfg_with_robot_initial_pose(self, scene_config: Any, pose: Pose) -> Any:
-        # ``pose`` is already lifted by the stand-height offset (see __init__ / set_initial_pose), so the
-        # base implementation sets the robot base as-is; we only add the stand placement here.
+        # ``pose`` is already lifted by the stand-height offset (see __init__ / set_initial_pose).
         scene_config = super()._update_scene_cfg_with_robot_initial_pose(scene_config, pose)
-        scene_config.stand.init_state.pos = pose.position_xyz
-        scene_config.stand.init_state.rot = pose.rotation_xyzw
+        stand_pose = self._stand_pose_from_robot_pose(pose)
+        scene_config.stand.init_state.pos = stand_pose.position_xyz
+        scene_config.stand.init_state.rot = stand_pose.rotation_xyzw
         return scene_config
 
     def set_initial_joint_pose(self, initial_joint_pose: list[float]) -> None:
@@ -267,7 +309,7 @@ class DroidSceneCfg:
     # TODO(alexmillane, 2025-07-28): We probably want to make the stand an optional addition.
     stand: AssetBaseCfg = AssetBaseCfg(
         prim_path="{ENV_REGEX_NS}/Robot_Stand",
-        init_state=AssetBaseCfg.InitialStateCfg(pos=[-0.05, 0.0, 0.0], rot=[0.0, 0.0, 0.0, 1.0]),
+        init_state=AssetBaseCfg.InitialStateCfg(pos=_STAND_ROOT_OFFSET_IN_ROBOT_FRAME, rot=[0.0, 0.0, 0.0, 1.0]),
         spawn=UsdFileCfg(
             usd_path=(
                 f"{ARENA_NUCLEUS_DIR}/Arena/assets/object_library/srl_robolab_assets/robots/franka_stand_grey.usda"