diff --git a/isaaclab_arena/assets/asset.py b/isaaclab_arena/assets/asset.py index 9cf4ccf0f4..81402c6a94 100644 --- a/isaaclab_arena/assets/asset.py +++ b/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 diff --git a/isaaclab_arena/assets/dummy_object.py b/isaaclab_arena/assets/dummy_object.py deleted file mode 100644 index 02f95daec7..0000000000 --- a/isaaclab_arena/assets/dummy_object.py +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright (c) 2025-2026, 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 __future__ import annotations - -import torch -import trimesh - -from isaaclab_arena.relations.collision_mode import CollisionMode -from isaaclab_arena.relations.relations import IsAnchor, Relation, RelationBase, UnaryRelation -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, quaternion_to_90_deg_z_quarters -from isaaclab_arena.utils.pose import Pose - - -class DummyObject: - """Dummy object for testing purposes without Isaac Sim dependencies.""" - - def __init__( - self, - name: str, - bounding_box: AxisAlignedBoundingBox, - initial_pose: Pose | None = None, - relations: list[RelationBase] = [], - collision_mesh: trimesh.Trimesh | None = None, - **kwargs, - ): - self.name = name - self.initial_pose = initial_pose - self.bounding_box = bounding_box - assert self.bounding_box is not None - self.relations = list(relations) - self._collision_mesh = collision_mesh - 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 add_relation(self, relation: RelationBase) -> None: - self.relations.append(relation) - - def get_relations(self) -> list[RelationBase]: - return 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 get_bounding_box(self) -> AxisAlignedBoundingBox: - """Get local bounding box (relative to object origin).""" - return self.bounding_box - - 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. - """ - if self.initial_pose is None: - return self.bounding_box - quarters = quaternion_to_90_deg_z_quarters(self.initial_pose.rotation_xyzw) - return self.bounding_box.rotated_90_around_z(quarters).translated(self.initial_pose.position_xyz) - - def get_corners_aabb(self, pos: torch.Tensor) -> torch.Tensor: - return self.bounding_box.get_corners_at(pos) - - def set_initial_pose(self, pose: Pose) -> None: - self.initial_pose = pose - - def get_initial_pose(self) -> Pose | None: - return self.initial_pose - - def is_initial_pose_set(self) -> bool: - return self.initial_pose is not None - - @property - def is_anchor(self) -> bool: - return any(isinstance(r, IsAnchor) for r in self.relations) - - def get_collision_mesh(self) -> trimesh.Trimesh | None: - """Return the collision mesh, or None to fall back to AABB.""" - return self._collision_mesh diff --git a/isaaclab_arena/assets/object.py b/isaaclab_arena/assets/object.py index 90ac4570cf..af3b53076c 100644 --- a/isaaclab_arena/assets/object.py +++ b/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,10 +63,6 @@ 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 @@ -77,22 +70,6 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox: 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: diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py index 19abe3949b..4ccd600b37 100644 --- a/isaaclab_arena/assets/object_base.py +++ b/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 diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index eaffa1c1eb..dae1c9a24d 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/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" diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index e55aa47a96..59ddfe2928 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -3,20 +3,27 @@ # # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + from collections.abc import Mapping -from typing import Any +from typing import TYPE_CHECKING, Any from isaaclab.envs import ManagerBasedRLMimicEnv +from isaaclab.managers import EventTermCfg from isaaclab.managers.recorder_manager import RecorderManagerBaseCfg -from isaaclab_arena.assets.asset import Asset from isaaclab_arena.embodiments.common.arm_mode import ArmMode +from isaaclab_arena.relations.placement_asset import PlacementAsset +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox 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 +from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange + +if TYPE_CHECKING: + import trimesh -class EmbodimentBase(Asset): +class EmbodimentBase(PlacementAsset): name: str | None = None tags: list[str] = ["embodiment"] @@ -49,17 +56,103 @@ def __init__( self.mimic_env: Any | None = None self.xr: Any | None = None self.termination_cfg: Any | None = None + self._collision_mesh: trimesh.Trimesh | None = None + """Lazily-extracted robot collision mesh, cached so the USD is opened once.""" + self.pose_event_cfg: EventTermCfg | None = None + """Root-pose reset event applied on environment reset; ``None`` when the embodiment has no configured pose.""" + + def get_bounding_box(self) -> AxisAlignedBoundingBox: + """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 + + assert self.scene_config is not None, "scene_config must be populated before placement" + robot = self.scene_config.robot + assert robot is not None, "scene_config.robot must be populated before placement" + spawn = robot.spawn + assert spawn.usd_path is not None, "scene_config.robot must use a USD spawn for placement" + scale = tuple(spawn.scale or (1.0, 1.0, 1.0)) + # TODO(zihaox): Account for configured initial joint positions in bounds and collision meshes. + return compute_local_bounding_box_from_usd(spawn.usd_path, scale) + + def get_collision_mesh(self) -> trimesh.Trimesh | None: + """Return the robot's collision mesh from its USD default prim, in the default joint pose.""" + if self._collision_mesh is None: + # Import locally because USD/pxr is available only after simulation initialization. + from pxr import Usd + + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_prim + + assert self.scene_config is not None, "scene_config must be populated before placement" + robot = self.scene_config.robot + assert robot is not None, "scene_config.robot must be populated before placement" + spawn = robot.spawn + assert spawn.usd_path is not None, "scene_config.robot must use a USD spawn for placement" + scale = tuple(spawn.scale or (1.0, 1.0, 1.0)) + stage = Usd.Stage.Open(spawn.usd_path) + assert stage is not None, f"could not open robot USD: {spawn.usd_path}" + default_prim = stage.GetDefaultPrim() or stage.GetPseudoRoot() + # The default prim scopes extraction to the robot, excluding sibling scene props + # (ground planes, stray objects) baked into some flattened articulation USDs. + self._collision_mesh = extract_trimesh_from_prim(stage, default_prim.GetPath().pathString, scale) + return self._collision_mesh + + def _get_initial_pose_as_pose(self) -> Pose | None: + """Return a single ``Pose`` for scene construction; ``PosePerEnv`` collapses to env 0.""" + initial_pose = self.initial_pose + if initial_pose is None: + return None + if isinstance(initial_pose, PosePerEnv): + return initial_pose.poses[0] + assert isinstance(initial_pose, Pose), "Embodiments support a fixed Pose or PosePerEnv only" + return initial_pose - def set_initial_pose(self, pose: Pose) -> None: + def _set_pose_state(self, pose: Pose | PosePerEnv) -> None: + """Store the configured pose; the construction pose is materialized in ``get_scene_cfg``.""" self.initial_pose = pose + def set_spawn_pose(self, pose: Pose) -> None: + """Set the scene-construction pose without rebuilding the reset event.""" + self._set_pose_state(pose) + + def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: + """Set the embodiment root pose and rebuild its pose reset event.""" + assert isinstance(pose, (Pose, PosePerEnv)), "Embodiments support a fixed Pose or PosePerEnv only" + self._set_pose_state(pose) + self.pose_event_cfg = self._init_pose_event_cfg() + + def _init_pose_event_cfg(self) -> EventTermCfg | None: + """Build the reset event that restores this embodiment's root pose (and auxiliary prims).""" + from isaaclab_arena.terms.events import reset_placement_asset_pose, reset_placement_asset_pose_per_env + + initial_pose = self.initial_pose + if initial_pose is None: + return None + if isinstance(initial_pose, PosePerEnv): + return EventTermCfg( + func=reset_placement_asset_pose_per_env, + mode="reset", + params={"write_pose_list": [self.layout_pose_to_scene_writes(pose) for pose in initial_pose.poses]}, + ) + assert isinstance(initial_pose, Pose), "Embodiments support a fixed Pose or PosePerEnv only" + return EventTermCfg( + func=reset_placement_asset_pose, + mode="reset", + params={"scene_writes": self.layout_pose_to_scene_writes(initial_pose)}, + ) + + def has_pose_reset_event(self) -> bool: + """Return whether the embodiment owns a root-pose reset event.""" + return self.pose_event_cfg is not None + def set_joint_initial_pos(self, joint_pos: Mapping[str, float]) -> None: """Update the robot's initial joint positions by joint name.""" - if self.scene_config is None or not hasattr(self.scene_config, "robot"): - raise RuntimeError("scene_config must be populated with a `robot` before calling `set_joint_initial_pos`.") - self.scene_config.robot.init_state.joint_pos.update(joint_pos) + assert self.scene_config is not None, "scene_config must be populated before setting joint positions" + robot = self.scene_config.robot + assert robot is not None, "scene_config.robot must be populated before setting joint positions" + robot.init_state.joint_pos.update(joint_pos) - def get_initial_pose(self) -> Pose: + def get_initial_pose(self) -> Pose | PosePerEnv: """Env-local robot base pose, resolved in order: the explicit ``initial_pose`` override if set, otherwise the ``scene_config`` robot ``init_state`` default.""" if self.initial_pose is not None: @@ -73,8 +166,9 @@ def get_initial_pose(self) -> Pose: ) def get_scene_cfg(self) -> Any: - if self.initial_pose is not None: - self.scene_config = self._update_scene_cfg_with_robot_initial_pose(self.scene_config, self.initial_pose) + construction_pose = self._get_initial_pose_as_pose() + if construction_pose is not None: + self.scene_config = self._update_scene_cfg_with_robot_initial_pose(self.scene_config, construction_pose) if self.enable_cameras: if self.camera_config is not None: return combine_configclass_instances( @@ -108,7 +202,16 @@ def get_commands_cfg(self) -> Any: return self.command_config def get_events_cfg(self) -> Any: - return self.event_config + if self.pose_event_cfg is None: + return self.event_config + from isaaclab_arena.utils.configclass import make_configclass + + pose_reset_cfg = make_configclass( + "EmbodimentPoseResetCfg", + [("robot_reset_pose", EventTermCfg, self.pose_event_cfg)], + )() + # Merge the pose reset last so it runs after joint/root resets in ``event_config``. + return combine_configclass_instances("EventsCfg", self.event_config, pose_reset_cfg) def get_mimic_env(self) -> ManagerBasedRLMimicEnv: return self.mimic_env @@ -138,10 +241,11 @@ def add_camera_variations(self, camera_rig: ArenaCameraCfg) -> None: self.add_variation(CameraIntrinsicsVariation(camera_name=camera_name, camera_rig=camera_rig)) def _update_scene_cfg_with_robot_initial_pose(self, scene_config: Any, pose: Pose) -> Any: - if scene_config is None or not hasattr(scene_config, "robot"): - raise RuntimeError("scene_config must be populated with a `robot` before calling `set_robot_initial_pose`.") - scene_config.robot.init_state.pos = pose.position_xyz - scene_config.robot.init_state.rot = pose.rotation_xyzw + assert scene_config is not None, "scene_config must be populated before setting the root pose" + robot = scene_config.robot + assert robot is not None, "scene_config.robot must be populated before setting the root pose" + robot.init_state.pos = pose.position_xyz + robot.init_state.rot = pose.rotation_xyzw return scene_config def get_recorder_term_cfg(self) -> RecorderManagerBaseCfg: @@ -150,7 +254,8 @@ def get_recorder_term_cfg(self) -> RecorderManagerBaseCfg: def get_termination_cfg(self) -> Any: return self.termination_cfg - def get_embodiment_name_in_scene(self) -> str: + def get_scene_name(self) -> str: + """Return the embodiment's Isaac Lab scene key.""" return "robot" def get_ee_frame_name(self, arm_mode: ArmMode) -> str: diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py index b5e2c306b3..a3add11dad 100644 --- a/isaaclab_arena/environments/arena_env_builder.py +++ b/isaaclab_arena/environments/arena_env_builder.py @@ -68,13 +68,12 @@ def __init__( self._placement_event_cfg: EventTermCfg | None = None def _solve_relations(self) -> None: - """Solve spatial relations for objects in the scene. + """Solve spatial relations for scene objects and the embodiment. This method: - 1. Collects all objects from the scene that have relations - 2. Builds an object-placement pool - 3. Reuses the object-only relation placer - 4. Applies solved positions either by writing fixed per-object initial poses + 1. Collects placement assets that have relations + 2. Builds a placement pool + 3. Applies solved positions either by writing fixed initial poses or by registering a pooled reset placement event Behaviour on reset depends on ``ObjectPlacerParams.resolve_on_reset``. @@ -83,10 +82,14 @@ def _solve_relations(self) -> None: * **True** (default) — registers a reset event that draws a fresh layout from the pool for each resetting environment. - * **False** — applies one layout per environment so per-object reset - events restore the same layout every time. + * **False** — assigns one fixed layout per environment. Object-only + scenes use per-object reset events; scenes with an embodiment use one + coordinated reset event. """ - objects_with_relations = self.arena_env.scene.get_objects_with_relations() + placement_assets = self.arena_env.scene.get_objects_with_relations() + embodiment = self.arena_env.embodiment + if embodiment is not None and embodiment.get_relations(): + placement_assets.append(embodiment) placer_params = self.arena_env.placer_params if placer_params is None: @@ -102,7 +105,7 @@ def _solve_relations(self) -> None: # TODO(xinjieyao, 2026-07-22): updated once robot-object co-placement is merged. placer_params.reachability_config.embodiment = self.arena_env.embodiment self._placement_event_cfg = solve_and_apply_relation_placement( - objects_with_relations, + placement_assets, num_envs=self.cfg.num_envs, placer_params=placer_params, scene_assets=self.arena_env.scene.assets.values(), diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py index 8895f2f919..d4b0535c93 100644 --- a/isaaclab_arena/environments/relation_solver_interface.py +++ b/isaaclab_arena/environments/relation_solver_interface.py @@ -11,19 +11,19 @@ from isaaclab_arena.relations.collision_mode import CollisionMode, get_object_collision_mode from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams -from isaaclab_arena.relations.placement_events import get_rotation_xyzw, solve_and_place_objects +from isaaclab_arena.relations.placement_events import get_pose_from_layout, solve_and_place_objects from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relations import get_anchor_objects -from isaaclab_arena.utils.pose import Pose, PosePerEnv -from isaaclab_arena.utils.yaw import rotate_quat_by_yaw, yaw_from_quat_xyzw +from isaaclab_arena.utils.pose import PosePerEnv if TYPE_CHECKING: from isaaclab.managers import EventTermCfg from isaaclab_arena.assets.asset import Asset - from isaaclab_arena.assets.object_base import ObjectBase from isaaclab_arena.assets.object_set import RigidObjectSet from isaaclab_arena.relations.collision_object import CollisionObject + from isaaclab_arena.relations.placement_asset import PlacementAsset + from isaaclab_arena.relations.placement_result import PlacementResult def _get_passive_collision_objects( @@ -37,16 +37,16 @@ def _get_passive_collision_objects( def solve_and_apply_relation_placement( - objects: list[ObjectBase], + assets: list[PlacementAsset], num_envs: int, placer_params: ObjectPlacerParams | None = None, collision_objects: list[CollisionObject] | None = None, scene_assets: Iterable[Asset | RigidObjectSet] | None = None, ) -> EventTermCfg | None: - """Solve relation placement and apply the result to object reset/static state. + """Solve relation placement and apply the result to asset reset/static state. Args: - objects: Objects with spatial predicates that should be relation-solved. + assets: Assets with spatial predicates that should be relation-solved. num_envs: Number of environments to prepare placements for. placer_params: Optional placement parameters. A shallow copy is used so this function can force pooled placement without mutating the caller's instance. @@ -59,10 +59,14 @@ def solve_and_apply_relation_placement( Reset event config to attach to the environment when placement should be resolved on reset. Returns ``None`` when no reset event is needed. """ - objects = list(objects) - if not objects: - print("No objects with relations found in scene. Skipping relation solving.") + assets = list(assets) + if not assets: + print("No assets with relations found in scene. Skipping relation solving.") return None + asset_names = {asset.name for asset in assets} + assert len(asset_names) == len(assets), "Placement asset names must be unique" + scene_keys = [asset.get_scene_name() for asset in assets] + assert len(set(scene_keys)) == len(scene_keys), "Placement assets map to duplicate scene keys" if placer_params is None: placer_params = ObjectPlacerParams() @@ -78,13 +82,11 @@ def solve_and_apply_relation_placement( collision_objects = _get_passive_collision_objects( scene_assets, include_background=_should_include_background_mesh( - objects, scene_assets, placer_params.solver_params.collision_mode + assets, scene_assets, placer_params.solver_params.collision_mode ), ) - # TODO(xinjieyao, 2026-05-22): Add joint object/embodiment placement once task-dependent - # reachability constraints are available. For now this always uses the object-only placer. placement_pool = PooledObjectPlacer( - objects=objects, + objects=assets, placer_params=placer_params, pool_size=num_envs * placer_params.min_unique_layouts_per_env, num_envs=num_envs, @@ -101,7 +103,7 @@ def solve_and_apply_relation_placement( ) return _apply_relation_placement_result( - objects=objects, + assets=assets, placer_params=placer_params, placement_pool=placement_pool, num_envs=num_envs, @@ -109,16 +111,16 @@ def solve_and_apply_relation_placement( def _should_include_background_mesh( - objects: list[ObjectBase], + assets: list[PlacementAsset], scene_assets: Iterable[Asset | RigidObjectSet], default_collision_mode: CollisionMode, ) -> bool: - """Return True when the default mode or any object/Background override resolves to MESH.""" + """Return True when the default mode or any relevant asset override resolves to MESH.""" from isaaclab_arena.assets.background import Background if default_collision_mode == CollisionMode.MESH: return True - if any(get_object_collision_mode(obj, default_collision_mode) == CollisionMode.MESH for obj in objects): + if any(get_object_collision_mode(asset, default_collision_mode) == CollisionMode.MESH for asset in assets): return True return any( isinstance(asset, Background) and get_object_collision_mode(asset, default_collision_mode) == CollisionMode.MESH @@ -127,111 +129,104 @@ def _should_include_background_mesh( def _apply_relation_placement_result( - objects: list[ObjectBase], + assets: list[PlacementAsset], placer_params: ObjectPlacerParams, placement_pool: PooledObjectPlacer, num_envs: int, ) -> EventTermCfg | None: - """Apply selected layouts to object spawn state and build reset event config.""" - anchor_objects_set = set(get_anchor_objects(objects)) - # Prevent external pose-reset events from conflicting with relation-solved objects. - _validate_no_conflicting_pose_reset_events(objects, anchor_objects_set) + """Apply selected layouts to asset spawn state and build reset event config.""" + anchor_assets = set(get_anchor_objects(assets)) + # Prevent external pose-reset events from conflicting with relation-solved assets. + _validate_no_conflicting_pose_reset_events(assets, anchor_assets) - # Anchor objects do not move, so no need to apply reset event. - if anchor_objects_set == set(objects): + # Anchor assets do not move, so no need to apply reset event. + if anchor_assets == set(assets): return None if placer_params.resolve_on_reset: return _apply_dynamic_spawn_pose( - objects=objects, + assets=assets, placement_pool=placement_pool, - anchor_objects_set=anchor_objects_set, + anchor_assets=anchor_assets, ) + # Every placement asset (objects and embodiments) stores its solved pose as a PosePerEnv and + # owns a per-asset reset event, so static layouts need no coordinated place-from-layouts event. _apply_static_initial_poses( - objects=objects, + assets=assets, placement_pool=placement_pool, - anchor_objects_set=anchor_objects_set, + anchor_assets=anchor_assets, num_envs=num_envs, ) + for asset in assets: + if asset in anchor_assets: + continue + assert asset.has_pose_reset_event(), ( + f"Static relation placement stored a per-env pose for non-anchor asset '{asset.name}', but it " + "owns no reset event, so its solved layout would be silently discarded on every reset." + ) return None def _apply_dynamic_spawn_pose( - objects: list[ObjectBase], + assets: list[PlacementAsset], placement_pool: PooledObjectPlacer, - anchor_objects_set: set[ObjectBase], + anchor_assets: set[PlacementAsset], ) -> EventTermCfg: """Set initial spawn pose from one layout and return the reset placement event.""" from isaaclab.managers import EventTermCfg - # For env-indexed pools this seeds from env 0; the first reset overwrites with per-env layouts. - layout = placement_pool.sample_with_replacement(1)[0] - for obj in objects: - if obj in anchor_objects_set: - continue - pos = layout.positions.get(obj) - if pos is None: - continue - base_rot = get_rotation_xyzw(obj) - marker_yaw = yaw_from_quat_xyzw(base_rot) - total_yaw = layout.orientations.get(obj, marker_yaw) - rot = rotate_quat_by_yaw(base_rot, total_yaw - marker_yaw) - object_cfg = getattr(obj, "object_cfg", None) - assert object_cfg is not None, f"Object '{obj.name}' must have object_cfg initialized before placement." - object_cfg.init_state.pos = pos - object_cfg.init_state.rot = rot + # Scene assets need a valid construction pose before reset events can run. + # This non-consuming env-0 sample is bootstrap-only; reset draws independently per env. + [construction_layout] = placement_pool.sample_with_replacement(1) + _seed_spawn_config_from_layout(assets, anchor_assets, construction_layout) return EventTermCfg( func=solve_and_place_objects, mode="reset", params={ - "objects": objects, + "assets": assets, "placement_pool": placement_pool, }, ) +def _seed_spawn_config_from_layout( + assets: list[PlacementAsset], + anchor_assets: set[PlacementAsset], + layout: PlacementResult, +) -> None: + """Write one solved layout into the single-pose scene configuration.""" + for asset in assets: + if asset in anchor_assets: + continue + pose = get_pose_from_layout(asset, layout) + asset.set_spawn_pose(pose) + + def _apply_static_initial_poses( - objects: list[ObjectBase], + assets: list[PlacementAsset], placement_pool: PooledObjectPlacer, - anchor_objects_set: set[ObjectBase], + anchor_assets: set[PlacementAsset], num_envs: int, ) -> None: """Apply fixed per-environment poses for ``resolve_on_reset=False``.""" layouts = placement_pool.sample_with_replacement(num_envs) - for obj in objects: - if obj in anchor_objects_set: + for asset in assets: + if asset in anchor_assets: continue - base_rotation_xyzw = get_rotation_xyzw(obj) - poses = [] - missing_envs: list[int] = [] - for env_idx in range(num_envs): - pos = layouts[env_idx].positions.get(obj) - if pos is None: - missing_envs.append(env_idx) - else: - marker_yaw = yaw_from_quat_xyzw(base_rotation_xyzw) - total_yaw = layouts[env_idx].orientations.get(obj, marker_yaw) - rotation_xyzw = rotate_quat_by_yaw(base_rotation_xyzw, total_yaw - marker_yaw) - poses.append(Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw)) - if missing_envs: - print( - f"Warning: Object '{obj.name}' is missing positions in {len(missing_envs)} env(s) " - f"(env ids: {missing_envs}); skipping set_initial_pose for this object." - ) - else: - obj.set_initial_pose(PosePerEnv(poses=poses)) + poses = [get_pose_from_layout(asset, layouts[env_idx]) for env_idx in range(num_envs)] + asset.set_initial_pose(PosePerEnv(poses=poses)) def _validate_no_conflicting_pose_reset_events( - objects: list[ObjectBase], - anchor_objects_set: set[ObjectBase], + assets: list[PlacementAsset], + anchor_assets: set[PlacementAsset], ) -> None: - """Reject conflicting explicit pose-reset events on relation-solved objects.""" - for obj in objects: - assert not (obj not in anchor_objects_set and getattr(obj, "event_cfg", None) is not None), ( - f"Non-anchor object '{obj.name}' has an explicit pose-reset event. " + """Reject conflicting explicit pose-reset events on relation-solved assets.""" + for asset in assets: + assert not (asset not in anchor_assets and asset.has_pose_reset_event()), ( + f"Non-anchor asset '{asset.name}' has an explicit pose-reset event. " "Relational solving should not be combined with explicit setting of " - "poses on non-anchor objects." + "poses on non-anchor assets." ) diff --git a/isaaclab_arena/relations/bounding_box_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py index 1428d52e5f..b7fc3586f6 100644 --- a/isaaclab_arena/relations/bounding_box_helpers.py +++ b/isaaclab_arena/relations/bounding_box_helpers.py @@ -5,7 +5,7 @@ """Bounding-box helpers for heterogeneous placement. -Keeps num_envs and per-env geometry logic out of ObjectBase. +Keeps num_envs and per-env geometry logic out of placement assets. """ from __future__ import annotations @@ -16,17 +16,17 @@ from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase + from isaaclab_arena.relations.placement_asset import PlacementAsset -def has_heterogeneous_objects(objects: list[ObjectBase]) -> bool: +def has_heterogeneous_objects(objects: list[PlacementAsset]) -> bool: """Return whether placement must use env-specific object geometry.""" from isaaclab_arena.assets.object_set import RigidObjectSet return any(isinstance(obj, RigidObjectSet) for obj in objects) -def assign_variants_for_envs(objects: list[ObjectBase], num_envs: int, placement_seed: int | None = None) -> None: +def assign_variants_for_envs(objects: list[PlacementAsset], num_envs: int, placement_seed: int | None = None) -> None: """Assign per-env variants on every RigidObjectSet in the list. Placers call this once they know the real environment count, before @@ -44,7 +44,7 @@ def assign_variants_for_envs(objects: list[ObjectBase], num_envs: int, placement variant_set_idx += 1 -def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBoundingBox: +def get_bounding_box_per_env(obj: PlacementAsset, num_envs: int) -> AxisAlignedBoundingBox: """Return bounding boxes expanded to (num_envs, 3). RigidObjectSet delegates to its own get_bounding_box_per_env. @@ -72,7 +72,7 @@ class PerEnvBoundingBoxes: (num_envs * candidates_per_env, 3), grouped contiguously by env. """ - object_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] + object_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox] num_envs: int def __post_init__(self) -> None: @@ -85,7 +85,7 @@ def __post_init__(self) -> None: 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[ObjectBase, AxisAlignedBoundingBox]: + def get_bounding_boxes_for_env_id(self, env_id: int) -> dict[PlacementAsset, AxisAlignedBoundingBox]: """Return object bboxes for a single env (each (1, 3)), used for per-env initialization and validation.""" return { obj: AxisAlignedBoundingBox( @@ -95,7 +95,7 @@ def get_bounding_boxes_for_env_id(self, env_id: int) -> dict[ObjectBase, AxisAli for obj, bbox in self.object_bboxes.items() } - def get_bounding_boxes_for_all_envs(self) -> list[dict[ObjectBase, AxisAlignedBoundingBox]]: + def get_bounding_boxes_for_all_envs(self) -> list[dict[PlacementAsset, AxisAlignedBoundingBox]]: """Return one-env bbox dicts for every env. The outer list has length num_envs. Each bbox has min_point/max_point @@ -105,7 +105,7 @@ def get_bounding_boxes_for_all_envs(self) -> list[dict[ObjectBase, AxisAlignedBo def get_bounding_boxes_for_solver_candidates( self, candidates_per_env: int - ) -> dict[ObjectBase, AxisAlignedBoundingBox]: + ) -> dict[PlacementAsset, AxisAlignedBoundingBox]: """Return bboxes tiled to one row per solver candidate. Each bbox has shape (num_envs * candidates_per_env, 3). Rows are grouped @@ -121,7 +121,7 @@ def get_bounding_boxes_for_solver_candidates( } -def build_per_env_bounding_boxes(objects: list[ObjectBase], num_envs: int) -> PerEnvBoundingBoxes: +def build_per_env_bounding_boxes(objects: list[PlacementAsset], 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 diff --git a/isaaclab_arena/relations/mesh_pair_cache.py b/isaaclab_arena/relations/mesh_pair_cache.py index 8871cae7c6..1adf8fac13 100644 --- a/isaaclab_arena/relations/mesh_pair_cache.py +++ b/isaaclab_arena/relations/mesh_pair_cache.py @@ -14,8 +14,8 @@ import warp as wp if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase from isaaclab_arena.relations.collision_object import CollisionObject + from isaaclab_arena.relations.placement_asset import PlacementAsset class MeshPairEntry(NamedTuple): @@ -24,10 +24,10 @@ class MeshPairEntry(NamedTuple): Dimensions: S = sphere count for this pair's subject, B = batch_size. """ - subject: ObjectBase + subject: PlacementAsset """Subject (sphere source) object.""" - obstacle: ObjectBase | CollisionObject + obstacle: PlacementAsset | CollisionObject """Obstacle (mesh target) object.""" obstacle_is_fixed: bool @@ -86,10 +86,10 @@ class MeshPairCache: all_radii: torch.Tensor """(S,) sphere radii, concatenated across pairs.""" - pair_subject_objs: list[ObjectBase] + pair_subject_objs: list[PlacementAsset] """(P,) subject (sphere source) object reference per pair.""" - pair_obstacle_objs: list[ObjectBase | CollisionObject] + pair_obstacle_objs: list[PlacementAsset | CollisionObject] """(P,) obstacle (mesh target) object reference per pair.""" pair_subject_applies_yaw: list[bool] diff --git a/isaaclab_arena/relations/no_overlap_aabb.py b/isaaclab_arena/relations/no_overlap_aabb.py index 8d941a55d8..4b9d4e6fb3 100644 --- a/isaaclab_arena/relations/no_overlap_aabb.py +++ b/isaaclab_arena/relations/no_overlap_aabb.py @@ -17,8 +17,8 @@ from isaaclab_arena.relations.relations import On if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase from isaaclab_arena.relations.collision_object import CollisionObject + from isaaclab_arena.relations.placement_asset import PlacementAsset from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache @@ -83,7 +83,7 @@ def compute_no_overlap_loss_aabb( on_pairs.add((id(obj), id(rel.parent))) on_pairs.add((id(rel.parent), id(obj))) - extents: dict[ObjectBase | CollisionObject, tuple[torch.Tensor, torch.Tensor]] = {} + extents: dict[PlacementAsset | CollisionObject, tuple[torch.Tensor, torch.Tensor]] = {} for obj in non_anchor_objects: pos = state.get_position(obj) bbox = state.get_bbox(obj) @@ -157,7 +157,7 @@ def compute_no_overlap_loss_aabb( def _fixed_pair_is_covered_by_mesh_collision( state: RelationSolverState, - subject: ObjectBase, + subject: PlacementAsset, obstacle: CollisionObject, mesh_manager: WarpMeshAndSphereCache, default_collision_mode: CollisionMode, @@ -175,8 +175,8 @@ def _fixed_pair_is_covered_by_mesh_collision( def _dynamic_pair_is_covered_by_mesh_collision( state: RelationSolverState, - a: ObjectBase, - b: ObjectBase, + a: PlacementAsset, + b: PlacementAsset, mesh_manager: WarpMeshAndSphereCache, default_collision_mode: CollisionMode, ) -> bool: @@ -192,7 +192,7 @@ def _dynamic_pair_is_covered_by_mesh_collision( def _has_mesh_or_invariant_bbox( state: RelationSolverState, - obj: ObjectBase, + obj: PlacementAsset, mesh_manager: WarpMeshAndSphereCache, default_collision_mode: CollisionMode, ) -> bool: diff --git a/isaaclab_arena/relations/no_overlap_mesh.py b/isaaclab_arena/relations/no_overlap_mesh.py index 69395fb4ee..e5fb52b8f6 100644 --- a/isaaclab_arena/relations/no_overlap_mesh.py +++ b/isaaclab_arena/relations/no_overlap_mesh.py @@ -22,8 +22,8 @@ from isaaclab_arena.utils.yaw import rotate_points_by_yaw_batch, yaw_from_quat_xyzw if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase from isaaclab_arena.relations.collision_object import CollisionObject + from isaaclab_arena.relations.placement_asset import PlacementAsset from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox @@ -32,7 +32,7 @@ def compute_no_overlap_loss_mesh( state: RelationSolverState, mesh_cache: MeshPairCache | None, mesh_manager: WarpMeshAndSphereCache, - orientations: list[dict[ObjectBase, float]] | None, + orientations: list[dict[PlacementAsset, float]] | None, clearance_m: float, slope: float, debug: bool, @@ -211,7 +211,7 @@ def _collect_mesh_pairs( state: RelationSolverState, manager: WarpMeshAndSphereCache, non_anchor_objects: list, - fixed_obstacles: list[ObjectBase | CollisionObject], + fixed_obstacles: list[PlacementAsset | CollisionObject], on_pairs: set[tuple[int, int]], device: torch.device, warned_no_mesh: set[str], @@ -361,7 +361,7 @@ def _collect_mesh_pairs( def _get_subject_spheres( mesh: trimesh.Trimesh | None, bbox: AxisAlignedBoundingBox, - obj: ObjectBase, + obj: PlacementAsset, manager: WarpMeshAndSphereCache, device: torch.device, ) -> torch.Tensor | None: diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py index e2108923c1..b6406a1ac3 100644 --- a/isaaclab_arena/relations/object_placer.py +++ b/isaaclab_arena/relations/object_placer.py @@ -29,8 +29,8 @@ from isaaclab_arena.utils.yaw import rotate_quat_by_yaw, wrap_angle_to_pi, yaw_from_quat_xyzw, yaw_toward_positions if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase from isaaclab_arena.relations.collision_object import CollisionObject + from isaaclab_arena.relations.placement_asset import PlacementAsset from isaaclab_arena.relations.placement_validators import PlacementValidator @@ -41,13 +41,13 @@ class PlacementCandidate: loss: float """Loss value returned by the solver.""" - positions: dict[ObjectBase, tuple[float, float, float]] + positions: dict[PlacementAsset, tuple[float, float, float]] """Solved positions for each object.""" validation_results: PlacementValidationResults """Per-check validation results for this candidate's layout.""" - orientations: dict[ObjectBase, float] = field(default_factory=dict) + orientations: dict[PlacementAsset, float] = field(default_factory=dict) """Placement-computed absolute world Z-yaws. Omitted objects retain their marker orientation.""" @property @@ -82,7 +82,7 @@ def __init__(self, params: ObjectPlacerParams | None = None): def place( self, - objects: list[ObjectBase], + objects: list[PlacementAsset], num_envs: int = 1, collision_objects: list[CollisionObject] | None = None, ) -> list[PlacementResult]: @@ -135,7 +135,7 @@ def place( def place_ranked_per_env( self, - objects: list[ObjectBase], + objects: list[PlacementAsset], num_envs: int, results_per_env: int, collision_objects: list[CollisionObject] | None = None, @@ -170,8 +170,8 @@ def place_ranked_per_env( def _prepare_placement( self, - objects: list[ObjectBase], - ) -> tuple[set[ObjectBase], torch.Generator | None]: + objects: list[PlacementAsset], + ) -> tuple[set[PlacementAsset], torch.Generator | None]: """Validate placement inputs and allocate an RNG seeded per candidate later.""" object_set = set(objects) for obj in objects: @@ -204,8 +204,8 @@ def _prepare_placement( def _place_ranked( self, - objects: list[ObjectBase], - anchor_objects_set: set[ObjectBase], + objects: list[PlacementAsset], + anchor_objects_set: set[PlacementAsset], num_envs: int, candidates_per_env: int, attempts_per_result: int, @@ -226,8 +226,8 @@ def _place_ranked( unrotated_candidate_bboxes = env_bboxes.get_bounding_boxes_for_solver_candidates(candidates_per_env) per_env_bboxes = env_bboxes.get_bounding_boxes_for_all_envs() - initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = [] - orientations_per_candidate: list[dict[ObjectBase, float]] = [] + initial_positions: list[dict[PlacementAsset, tuple[float, float, float]]] = [] + orientations_per_candidate: list[dict[PlacementAsset, float]] = [] for candidate_idx in range(num_candidates): cur_env = candidate_idx // candidates_per_env if generator is not None: @@ -332,11 +332,11 @@ def _print_ranked_summary( def _generate_initial_positions( self, - objects: list[ObjectBase], - anchor_objects: set[ObjectBase], - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + objects: list[PlacementAsset], + anchor_objects: set[PlacementAsset], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], generator: torch.Generator | None = None, - ) -> dict[ObjectBase, tuple[float, float, float]]: + ) -> dict[PlacementAsset, tuple[float, float, float]]: """Generate initial positions for all objects. Anchors keep their initial_pose. Objects with an On relation are initialized within @@ -356,7 +356,7 @@ def _generate_initial_positions( cx, cy, cz = float(anchor_bbox.center[0, 0]), float(anchor_bbox.center[0, 1]), float(anchor_bbox.center[0, 2]) - positions: dict[ObjectBase, tuple[float, float, float]] = {} + positions: dict[PlacementAsset, tuple[float, float, float]] = {} for obj in objects: if obj in anchor_objects: initial_pose = obj.get_initial_pose() @@ -375,8 +375,8 @@ def _generate_initial_positions( @staticmethod def _get_world_bbox_for_init( - obj: ObjectBase, - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + obj: PlacementAsset, + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], ) -> AxisAlignedBoundingBox: initial_pose = obj.get_initial_pose() assert isinstance( @@ -386,17 +386,17 @@ def _get_world_bbox_for_init( def _generate_initial_orientations( self, - objects: list[ObjectBase], - anchor_objects: set[ObjectBase], + objects: list[PlacementAsset], + anchor_objects: set[PlacementAsset], generator: torch.Generator | None = None, - ) -> dict[ObjectBase, float]: + ) -> dict[PlacementAsset, float]: """Sample absolute world Z-yaws for non-anchor objects without FaceTo. Marker yaw is included; random_yaw_init adds a sampled delta. Roll/pitch marker objects are omitted so their requested rotation is applied verbatim; their footprint is enclosed by _rotate_candidate_bboxes so overlap validation stays sound. """ - orientations: dict[ObjectBase, float] = {} + orientations: dict[PlacementAsset, float] = {} for obj in objects: marker = get_relation(obj, RotateAroundSolution) has_roll_pitch = marker is not None and (marker.roll_rad != 0.0 or marker.pitch_rad != 0.0) @@ -416,8 +416,8 @@ def _generate_initial_orientations( @staticmethod def _apply_face_to_orientations( - positions_per_candidate: list[dict[ObjectBase, tuple[float, float, float]]], - orientations_per_candidate: list[dict[ObjectBase, float]], + positions_per_candidate: list[dict[PlacementAsset, tuple[float, float, float]]], + orientations_per_candidate: list[dict[PlacementAsset, float]], ) -> None: """Write defined FaceTo yaws into each candidate's orientation dictionary in place. @@ -439,10 +439,10 @@ def _apply_face_to_orientations( @staticmethod def _rotate_candidate_bboxes( - objects: list[ObjectBase], - candidate_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], - orientations_per_candidate: list[dict[ObjectBase, float]], - ) -> dict[ObjectBase, AxisAlignedBoundingBox]: + objects: list[PlacementAsset], + candidate_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], + orientations_per_candidate: list[dict[PlacementAsset, float]], + ) -> dict[PlacementAsset, AxisAlignedBoundingBox]: """Replace each candidate's bbox with the AABB enclosing its fully-oriented object. Composes each object's static RotateAroundSolution marker rotation (roll/pitch/yaw) with the @@ -452,7 +452,7 @@ def _rotate_candidate_bboxes( no rotation are returned unchanged, keeping the no-rotation path exact. """ num_candidates = len(orientations_per_candidate) - rotated: dict[ObjectBase, AxisAlignedBoundingBox] = {} + rotated: dict[PlacementAsset, AxisAlignedBoundingBox] = {} for obj in objects: bbox = candidate_bboxes[obj] marker = get_relation(obj, RotateAroundSolution) @@ -473,18 +473,18 @@ def _rotate_candidate_bboxes( @staticmethod def _get_bounding_boxes_for_candidate_index( - bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], candidate_idx: int, - ) -> dict[ObjectBase, AxisAlignedBoundingBox]: + ) -> dict[PlacementAsset, AxisAlignedBoundingBox]: """Slice one candidate's bboxes (each (1, 3)) out of the stacked (num_candidates, 3) boxes.""" return {obj: bbox[candidate_idx] for obj, bbox in bboxes.items()} def _get_on_parent_world_bbox( self, - parent: ObjectBase, - anchor_objects: set[ObjectBase], + parent: PlacementAsset, + anchor_objects: set[PlacementAsset], anchor_bbox: AxisAlignedBoundingBox, - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], ) -> AxisAlignedBoundingBox: """Resolve the world bbox of an On relation's parent for initialization purposes. @@ -504,10 +504,10 @@ def _get_on_parent_world_bbox( def _compute_on_guided_position( self, - obj: ObjectBase, - anchor_objects: set[ObjectBase], + obj: PlacementAsset, + anchor_objects: set[PlacementAsset], anchor_bbox: AxisAlignedBoundingBox, - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], generator: torch.Generator | None = None, ) -> tuple[float, float, float]: """Compute an initial position for an object with an On relation. @@ -576,9 +576,9 @@ def _sample_axis_position( def _validate_candidates( self, - positions: list[dict[ObjectBase, tuple[float, float, float]]], - orientations: list[dict[ObjectBase, float]], - bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]], + positions: list[dict[PlacementAsset, tuple[float, float, float]]], + orientations: list[dict[PlacementAsset, float]], + bboxes: list[dict[PlacementAsset, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[PlacementValidationResults]: """Run every enabled validator over all candidates and collect per-candidate results. @@ -634,9 +634,9 @@ def _validate_candidates( def _run_inexpensive_checks( self, - positions: list[dict[ObjectBase, tuple[float, float, float]]], - orientations: list[dict[ObjectBase, float]], - bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]], + positions: list[dict[PlacementAsset, tuple[float, float, float]]], + orientations: list[dict[PlacementAsset, float]], + bboxes: list[dict[PlacementAsset, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], layout_pass_verdicts_by_check: dict[str, list[bool]], num_layouts_evaluated_by_check: dict[str, int], @@ -652,9 +652,9 @@ def _run_inexpensive_checks( def _run_expensive_checks( self, - positions: list[dict[ObjectBase, tuple[float, float, float]]], - orientations: list[dict[ObjectBase, float]], - bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]], + positions: list[dict[PlacementAsset, tuple[float, float, float]]], + orientations: list[dict[PlacementAsset, float]], + bboxes: list[dict[PlacementAsset, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], required: set[str] | None, layout_pass_verdicts_by_check: dict[str, list[bool]], @@ -700,9 +700,9 @@ def _passes_required_checks( def _apply_poses( self, - positions_per_env: list[dict[ObjectBase, tuple[float, float, float]]], - anchor_objects: set[ObjectBase], - orientations_per_env: list[dict[ObjectBase, float]], + positions_per_env: list[dict[PlacementAsset, tuple[float, float, float]]], + anchor_objects: set[PlacementAsset], + orientations_per_env: list[dict[PlacementAsset, float]], ) -> None: """Apply solved positions and orientations to non-anchor objects. diff --git a/isaaclab_arena/relations/placement_asset.py b/isaaclab_arena/relations/placement_asset.py new file mode 100644 index 0000000000..e6c188f361 --- /dev/null +++ b/isaaclab_arena/relations/placement_asset.py @@ -0,0 +1,99 @@ +# Copyright (c) 2026, 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 + +"""Shared model for assets whose poses are relation-solved.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from isaaclab_arena.assets.asset import Asset +from isaaclab_arena.relations.collision_mode import CollisionMode +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 + +if TYPE_CHECKING: + import trimesh + + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + + +class PlacementAsset(Asset, ABC): + """Asset whose root pose can be constrained by spatial relations.""" + + def __init__(self, name: str, tags: list[str] | None = None, **kwargs) -> None: + super().__init__(name=name, tags=tags, **kwargs) + self.initial_pose: Pose | PoseRange | PosePerEnv | None = None + self.relations: list[RelationBase] = [] + # None delegates collision-mode selection to the solver. + self.collision_mode: CollisionMode | None = None + # Whether to replace a non-watertight collision mesh with its convex hull. + self.repair_collision_mesh_non_watertight = True + + def add_relation(self, relation: RelationBase) -> None: + """Attach a relation to the asset.""" + self.relations.append(relation) + + def get_relations(self) -> list[RelationBase]: + """Return all relations attached to the asset.""" + return self.relations + + def get_spatial_relations(self) -> list[RelationBase]: + """Return spatial constraints, excluding placement markers.""" + return [relation for relation in self.relations if isinstance(relation, (Relation, UnaryRelation))] + + @property + def is_anchor(self) -> bool: + """Return whether the asset is fixed during relation solving.""" + return any(isinstance(relation, IsAnchor) for relation in self.relations) + + def get_initial_pose(self) -> Pose | PoseRange | PosePerEnv | None: + """Return the configured root pose.""" + return self.initial_pose + + def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: + """Set the configured root pose. + + Accepts a fixed ``Pose``, a ``PoseRange`` (randomized on reset), or a ``PosePerEnv`` + (a distinct pose per environment); the interpretation is subclass-specific. + """ + self.initial_pose = pose + + def set_spawn_pose(self, pose: Pose) -> None: + """Set the root pose used when constructing the scene.""" + self.set_initial_pose(pose) + + def layout_pose_to_scene_writes(self, layout_pose: Pose) -> list[tuple[str, Pose]]: + """Return the ``(scene entity name, env-local pose)`` writes that realize a solved layout pose. + + A simple asset places only its own root; a compound asset (e.g. a robot on a separate + stand) overrides this to also place auxiliary prims that move with the root. + """ + return [(self.get_scene_name(), layout_pose)] + + def has_pose_reset_event(self) -> bool: + """Return whether the asset owns a root-pose reset event.""" + return False + + @abstractmethod + def get_bounding_box(self) -> AxisAlignedBoundingBox: + """Return root-relative axis-aligned bounds.""" + + def get_world_bounding_box(self) -> AxisAlignedBoundingBox: + """Return bounds transformed by a fixed root pose with a quarter-turn Z rotation. + + Unset, ranged, and per-environment poses leave the root-relative bounds unchanged. + """ + bounding_box = self.get_bounding_box() + initial_pose = self.get_initial_pose() + if not isinstance(initial_pose, Pose): + return bounding_box + quarters = quaternion_to_90_deg_z_quarters(initial_pose.rotation_xyzw) + return bounding_box.rotated_90_around_z(quarters).translated(initial_pose.position_xyz) + + def get_collision_mesh(self) -> trimesh.Trimesh | None: + """Return this asset's collision mesh, or ``None`` if it has none.""" diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py index 7a4e330ec1..8b85d5c4cd 100644 --- a/isaaclab_arena/relations/placement_events.py +++ b/isaaclab_arena/relations/placement_events.py @@ -10,13 +10,13 @@ from isaaclab_arena.relations.relations import RotateAroundSolution, get_anchor_objects from isaaclab_arena.utils.pose import Pose -from isaaclab_arena.utils.velocity import Velocity +from isaaclab_arena.utils.scene_pose_writes import write_scene_root_poses_to_sim from isaaclab_arena.utils.yaw import rotate_quat_by_yaw, yaw_from_quat_xyzw if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv - from isaaclab_arena.assets.object_base import ObjectBase + from isaaclab_arena.relations.placement_asset import PlacementAsset from isaaclab_arena.relations.placement_result import PlacementResult from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer @@ -42,64 +42,93 @@ def get_placement_pool(env) -> PooledObjectPlacer | None: return term_cfg.params.get("placement_pool") -def get_rotation_xyzw(obj: ObjectBase) -> tuple[float, float, float, float]: - """Return the RotateAroundSolution rotation for *obj*, or identity if none.""" - rotate_marker = next((r for r in obj.get_relations() if isinstance(r, RotateAroundSolution)), None) +def get_rotation_xyzw(asset: PlacementAsset) -> tuple[float, float, float, float]: + """Return the RotateAroundSolution rotation for an asset, or identity if none.""" + rotate_marker = next((r for r in asset.get_relations() if isinstance(r, RotateAroundSolution)), None) return rotate_marker.get_rotation_xyzw() if rotate_marker else IDENTITY_ROTATION_XYZW -def get_base_rotation_per_object(objects: list[ObjectBase]) -> dict[ObjectBase, tuple[float, float, float, float]]: - """Return the base rotation for each object.""" - return {obj: get_rotation_xyzw(obj) for obj in objects} +def get_base_rotation_per_asset( + assets: list[PlacementAsset], +) -> dict[PlacementAsset, tuple[float, float, float, float]]: + """Return the base rotation for each asset.""" + return {asset: get_rotation_xyzw(asset) for asset in assets} -def get_movable_object_names( - objects: list[ObjectBase], - anchor_objects_set: set[ObjectBase], +def get_pose_from_layout(asset: PlacementAsset, layout: PlacementResult) -> Pose: + """Return an asset pose from a solved layout.""" + assert asset in layout.positions, f"Placement layout is missing non-anchor asset '{asset.name}'" + base_rotation = get_rotation_xyzw(asset) + marker_yaw = yaw_from_quat_xyzw(base_rotation) + total_yaw = layout.orientations.get(asset, marker_yaw) + rotation = rotate_quat_by_yaw(base_rotation, total_yaw - marker_yaw) + return Pose(position_xyz=layout.positions[asset], rotation_xyzw=rotation) + + +def get_movable_asset_names( + assets: list[PlacementAsset], + anchor_assets: set[PlacementAsset], ) -> list[str]: - """Return the names of non-anchor objects.""" - return [obj.name for obj in objects if obj not in anchor_objects_set] + """Return scene names for non-anchor placement assets.""" + return [asset.get_scene_name() for asset in assets if asset not in anchor_assets] + + +def _write_scene_root_pose_to_sim( + env: ManagerBasedEnv, + scene_name: str, + sim_pose: Pose, + env_id: int, + env_id_tensor: torch.Tensor, +) -> None: + """Write one scene root pose for a single environment instance.""" + scene_asset = env.scene[scene_name] + pose_tensor = sim_pose.to_tensor(device=env.device).unsqueeze(0) + pose_tensor[0, :3] += env.scene.env_origins[env_id, :] + write_scene_root_poses_to_sim(scene_asset, scene_name, pose_tensor, env_id_tensor, env.device) def write_layout_to_sim( env: ManagerBasedEnv, env_id: int, result: PlacementResult, - anchor_objects_set: set[ObjectBase], - base_rotations: dict[ObjectBase, tuple[float, float, float, float]], + anchor_assets: set[PlacementAsset], + base_rotations: dict[PlacementAsset, tuple[float, float, float, float]], ) -> None: """Write one env's solved layout into the sim. Even writing zero velocity, the sim will still apply gravity and other forces from collisions, - so collided objects will still be subject to move. + so collided assets will still be subject to move. Args: env: The Isaac Lab ManagerBasedEnv environment. env_id: The environment index. result: The placement result to write to the sim. - anchor_objects_set: The set of anchor objects. - base_rotations: The base rotations for all objects. + anchor_assets: The set of anchor assets. + base_rotations: The base rotations for all assets. """ env_id_tensor = torch.tensor([env_id], device=env.device) - zero_velocity = Velocity.zero().to_tensor(device=env.device).unsqueeze(0) - for obj, pos in result.positions.items(): - if obj in anchor_objects_set: + missing_assets = [ + asset.name for asset in base_rotations if asset not in anchor_assets and asset not in result.positions + ] + assert not missing_assets, f"Placement layout is missing non-anchor assets: {missing_assets}" + for asset in result.positions: + if asset in anchor_assets: continue - asset = env.scene[obj.name] - marker_yaw = yaw_from_quat_xyzw(base_rotations[obj]) - total_yaw = result.orientations.get(obj, marker_yaw) - rotation_xyzw = rotate_quat_by_yaw(base_rotations[obj], total_yaw - marker_yaw) - pose = Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw) - pose_t_xyz_q_xyzw = pose.to_tensor(device=env.device).unsqueeze(0) - pose_t_xyz_q_xyzw[0, :3] += env.scene.env_origins[env_id, :] - asset.write_root_pose_to_sim(pose_t_xyz_q_xyzw, env_ids=env_id_tensor) - asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_id_tensor) + layout_pose = get_pose_from_layout(asset, result) + for scene_name, sim_pose in asset.layout_pose_to_scene_writes(layout_pose): + _write_scene_root_pose_to_sim( + env, + scene_name, + sim_pose, + env_id, + env_id_tensor, + ) def solve_and_place_objects( env: ManagerBasedEnv, env_ids: torch.Tensor | None, - objects: list[ObjectBase], + assets: list[PlacementAsset], placement_pool: PooledObjectPlacer, ) -> None: """Coordinated reset event that draws layouts from the pool and writes poses. @@ -111,7 +140,7 @@ def solve_and_place_objects( Args: env: The Isaac Lab environment. env_ids: 1-D tensor of environment indices being reset. - objects: Objects participating in relation solving. + assets: Assets participating in relation solving. placement_pool: Runtime pool of solved placement layouts. """ if env_ids is None or len(env_ids) == 0: @@ -122,8 +151,8 @@ def solve_and_place_objects( placement_pool.num_envs == num_scene_envs ), f"Placement pool has {placement_pool.num_envs} envs, but scene has {num_scene_envs} env origins." results_by_env = placement_pool.sample_for_envs(reset_env_ids) - anchor_objects_set = set(get_anchor_objects(objects)) - base_rotations = get_base_rotation_per_object(objects) + anchor_assets = set(get_anchor_objects(assets)) + base_rotations = get_base_rotation_per_asset(assets) for cur_env in reset_env_ids: result = results_by_env[cur_env] @@ -132,5 +161,5 @@ def solve_and_place_objects( "Warning: Writing best-loss fallback placement for " f"env {cur_env}; failed checks: {result.validation_results.get_failed_validation_check_names}." ) - # only write the non-anchor objects to the sim - write_layout_to_sim(env, cur_env, result, anchor_objects_set, base_rotations) + # Only write non-anchor assets to the sim. + write_layout_to_sim(env, cur_env, result, anchor_assets, base_rotations) diff --git a/isaaclab_arena/relations/placement_pool_validation.py b/isaaclab_arena/relations/placement_pool_validation.py index ee8f73aeb3..98bcb4fb9f 100644 --- a/isaaclab_arena/relations/placement_pool_validation.py +++ b/isaaclab_arena/relations/placement_pool_validation.py @@ -9,8 +9,8 @@ from isaaclab_arena.relations.physics_settle_params import PhysicsSettleParams from isaaclab_arena.relations.placement_events import ( - get_base_rotation_per_object, - get_movable_object_names, + get_base_rotation_per_asset, + get_movable_asset_names, get_placement_pool, write_layout_to_sim, ) @@ -21,7 +21,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv - from isaaclab_arena.relations.object_base import ObjectBase + from isaaclab_arena.relations.placement_asset import PlacementAsset from isaaclab_arena.relations.placement_result import PlacementResult from isaaclab_arena.relations.placement_validation import PlacementValidationResults from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer @@ -32,8 +32,8 @@ def _write_layout_to_envs_for_episode_index( layouts_per_env: list[list[PlacementResult]], num_envs: int, episode_index: int, - anchor_objects_set: set, - base_rotations: dict[ObjectBase, tuple[float, float, float, float]], + anchor_assets: set, + base_rotations: dict[PlacementAsset, tuple[float, float, float, float]], ) -> list[tuple[int, PlacementResult]]: """Write one layout per env for this episode; return the ``(env_id, layout)`` layouts written. @@ -45,7 +45,13 @@ def _write_layout_to_envs_for_episode_index( layouts = layouts_per_env[env_id] if episode_index < len(layouts): layout = layouts[episode_index] - write_layout_to_sim(env.unwrapped, env_id, layout, anchor_objects_set, base_rotations) + write_layout_to_sim( + env.unwrapped, + env_id, + layout, + anchor_assets, + base_rotations, + ) layouts_written.append((env_id, layout)) return layouts_written @@ -97,7 +103,6 @@ def validate_pool_layouts( ``(env_id, episode_index, checklist)`` for every layout, in ``(env_id, episode_index)`` order, or ``None`` when ``placement_pool`` is omitted and the env has no pooled layouts. """ - # No-ops when no layouts are stored in the pool if placement_pool is None: placement_pool = get_placement_pool(env) if placement_pool is None: @@ -105,10 +110,10 @@ def validate_pool_layouts( if settle_params is None: settle_params = PhysicsSettleParams() - objects = placement_pool.objects - anchor_objects_set = set(get_anchor_objects(objects)) - base_rotations = get_base_rotation_per_object(objects) - movable_object_names = get_movable_object_names(objects, anchor_objects_set) + assets = placement_pool.objects + anchor_assets = set(get_anchor_objects(assets)) + base_rotations = get_base_rotation_per_asset(assets) + movable_object_names = get_movable_asset_names(assets, anchor_assets) # The length of each env queue is controlled by min_unique_layouts_per_env in ObjectPlacerParams. layouts_per_env = placement_pool.layouts_per_env() @@ -127,7 +132,12 @@ def validate_pool_layouts( for episode_index in range(max_episodes): # Set layout, then settle and collect results in parallel. layouts = _write_layout_to_envs_for_episode_index( - env, layouts_per_env, num_envs, episode_index, anchor_objects_set, base_rotations + env, + layouts_per_env, + num_envs, + episode_index, + anchor_assets, + base_rotations, ) if layouts: physics_settle.step_physics(env, num_physics_steps, render=render) diff --git a/isaaclab_arena/relations/placement_result.py b/isaaclab_arena/relations/placement_result.py index 03a50e8a5b..670ade5a1f 100644 --- a/isaaclab_arena/relations/placement_result.py +++ b/isaaclab_arena/relations/placement_result.py @@ -9,19 +9,19 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase + from isaaclab_arena.relations.placement_asset import PlacementAsset from isaaclab_arena.relations.placement_validation import PlacementValidationResults @dataclass class PlacementResult: - """Solved object layout for one environment.""" + """Solved asset layout for one environment.""" validation_results: PlacementValidationResults """Validation checklist for the placement.""" - positions: dict[ObjectBase, tuple[float, float, float]] - """Final positions for each object.""" + positions: dict[PlacementAsset, tuple[float, float, float]] + """Final ``(x, y, z)`` positions in metres in the environment-local frame.""" final_loss: float """Loss value of the final placement.""" @@ -29,13 +29,10 @@ class PlacementResult: attempts: int """Number of attempts made.""" - orientations: dict[ObjectBase, float] = field(default_factory=dict) - """Placement-computed absolute world Z-yaws. Omitted objects retain their marker orientation.""" + orientations: dict[PlacementAsset, float] = field(default_factory=dict) + """Sparse map of world yaw angles ``theta_z`` in radians; omitted assets retain marker orientation.""" @property def success(self) -> bool: - """True when all required validation checks pass. - - place() returns a best-loss fallback even on failure; check this to tell validated from fallback. - """ + """Whether all required validation checks passed.""" return self.validation_results.do_all_required_validation_checks_pass() diff --git a/isaaclab_arena/relations/placement_validators.py b/isaaclab_arena/relations/placement_validators.py index 13b3fe8905..2832ad1657 100644 --- a/isaaclab_arena/relations/placement_validators.py +++ b/isaaclab_arena/relations/placement_validators.py @@ -26,9 +26,9 @@ from isaaclab_arena.utils.yaw import centers_in_target_frame, yaw_from_quat_xyzw, yaw_toward_positions if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase from isaaclab_arena.relations.collision_object import CollisionObject from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams + from isaaclab_arena.relations.placement_asset import PlacementAsset from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox @@ -62,9 +62,9 @@ def is_available(cls, params: ObjectPlacerParams) -> bool: @abstractmethod def validate_batch( self, - positions: list[dict[ObjectBase, tuple[float, float, float]]], - orientations: list[dict[ObjectBase, float]], - bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]], + positions: list[dict[PlacementAsset, tuple[float, float, float]]], + orientations: list[dict[PlacementAsset, float]], + bboxes: list[dict[PlacementAsset, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: """Return one pass/fail verdict per candidate layout. @@ -116,17 +116,17 @@ class OnRelationValidator(PlacementValidator): def validate_batch( self, - positions: list[dict[ObjectBase, tuple[float, float, float]]], - orientations: list[dict[ObjectBase, float]], - bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]], + positions: list[dict[PlacementAsset, tuple[float, float, float]]], + orientations: list[dict[PlacementAsset, float]], + bboxes: list[dict[PlacementAsset, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: return [self._validate(positions[i], bboxes[i]) for i in range(len(positions))] def _validate( self, - positions: dict[ObjectBase, tuple[float, float, float]], - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + positions: dict[PlacementAsset, tuple[float, float, float]], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], ) -> bool: """Validate each On relation; keep in sync with OnLossStrategy in relation_loss_strategies.py. @@ -201,17 +201,17 @@ class NextToValidator(PlacementValidator): def validate_batch( self, - positions: list[dict[ObjectBase, tuple[float, float, float]]], - orientations: list[dict[ObjectBase, float]], - bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]], + positions: list[dict[PlacementAsset, tuple[float, float, float]]], + orientations: list[dict[PlacementAsset, float]], + bboxes: list[dict[PlacementAsset, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: return [self._validate(positions[i], bboxes[i]) for i in range(len(positions))] def _validate( self, - positions: dict[ObjectBase, tuple[float, float, float]], - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + positions: dict[PlacementAsset, tuple[float, float, float]], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], ) -> bool: """Validate each NextTo relation: child on the requested side, facing edge within the relation's tolerance_m of distance_m from the parent edge. Shares next_to_violations with @@ -253,17 +253,17 @@ class NotNextToValidator(PlacementValidator): def validate_batch( self, - positions: list[dict[ObjectBase, tuple[float, float, float]]], - orientations: list[dict[ObjectBase, float]], - bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]], + positions: list[dict[PlacementAsset, tuple[float, float, float]]], + orientations: list[dict[PlacementAsset, float]], + bboxes: list[dict[PlacementAsset, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: return [self._validate(positions[i], bboxes[i]) for i in range(len(positions))] def _validate( self, - positions: dict[ObjectBase, tuple[float, float, float]], - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + positions: dict[PlacementAsset, tuple[float, float, float]], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], ) -> bool: """Validate each NotNextTo relation: child has cleared the keep-out zone beside the parent (within the relation's tolerance_m) via either route — back over the edge or past the @@ -314,17 +314,17 @@ class FaceToValidator(PlacementValidator): def validate_batch( self, - positions: list[dict[ObjectBase, tuple[float, float, float]]], - orientations: list[dict[ObjectBase, float]], - bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]], + positions: list[dict[PlacementAsset, tuple[float, float, float]]], + orientations: list[dict[PlacementAsset, float]], + bboxes: list[dict[PlacementAsset, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: return [self._validate(positions[i], orientations[i]) for i in range(len(positions))] def _validate( self, - positions: dict[ObjectBase, tuple[float, float, float]], - orientations: dict[ObjectBase, float] | None, + positions: dict[PlacementAsset, tuple[float, float, float]], + orientations: dict[PlacementAsset, float] | None, ) -> bool: """Validate that every FaceTo subject has a defined direction and computed yaw.""" for obj in positions: @@ -361,9 +361,9 @@ def __init__(self, params: ObjectPlacerParams) -> None: def validate_batch( self, - positions: list[dict[ObjectBase, tuple[float, float, float]]], - orientations: list[dict[ObjectBase, float]], - bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]], + positions: list[dict[PlacementAsset, tuple[float, float, float]]], + orientations: list[dict[PlacementAsset, float]], + bboxes: list[dict[PlacementAsset, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[bool]: return [ @@ -372,9 +372,9 @@ def validate_batch( def _validate( self, - positions: dict[ObjectBase, tuple[float, float, float]], - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], - orientations: dict[ObjectBase, float] | None, + positions: dict[PlacementAsset, tuple[float, float, float]], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], + orientations: dict[PlacementAsset, float] | None, collision_objects: list[CollisionObject] | None, ) -> bool: """AABB overlap check, falling through to mesh penetration for mesh-collision objects.""" @@ -391,7 +391,7 @@ def _validate( def _should_validate_mesh( self, - positions: dict[ObjectBase, tuple[float, float, float]], + positions: dict[PlacementAsset, tuple[float, float, float]], collision_objects: list[CollisionObject] | None, ) -> bool: """Return True when any object in this validation uses mesh collision.""" @@ -403,7 +403,7 @@ def _should_validate_mesh( @staticmethod def _collect_skip_pairs( - positions: dict[ObjectBase, tuple[float, float, float]], + positions: dict[PlacementAsset, tuple[float, float, float]], ) -> tuple[set[tuple], set[int]]: """Build On-pair skip set and anchor ID set from positioned objects. @@ -424,9 +424,9 @@ def _collect_skip_pairs( def _non_skip_pairs( self, - positions: dict[ObjectBase, tuple[float, float, float]], + positions: dict[PlacementAsset, tuple[float, float, float]], skip_mesh_pairs: bool = False, - ) -> Iterator[tuple[ObjectBase, ObjectBase]]: + ) -> Iterator[tuple[PlacementAsset, PlacementAsset]]: """Yield non-relation object pairs, optionally skipping pairs handled by mesh collision.""" on_pairs, anchor_ids = self._collect_skip_pairs(positions) mesh_manager = self._get_cpu_mesh_manager() if skip_mesh_pairs else None @@ -454,8 +454,8 @@ def _non_skip_pairs( def _validate_no_overlap( self, - positions: dict[ObjectBase, tuple[float, float, float]], - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + positions: dict[PlacementAsset, tuple[float, float, float]], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], collision_objects: list[CollisionObject] | None = None, skip_mesh_pairs: bool = False, ) -> bool: @@ -506,9 +506,9 @@ def _get_cpu_mesh_manager(self) -> WarpMeshAndSphereCache: def _validate_no_overlap_mesh( self, - positions: dict[ObjectBase, tuple[float, float, float]], - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], - orientations: dict[ObjectBase, float] | None = None, + positions: dict[PlacementAsset, tuple[float, float, float]], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], + orientations: dict[PlacementAsset, float] | None = None, collision_objects: list[CollisionObject] | None = None, ) -> bool: """Sphere-to-SDF overlap check; both-meshless pairs fall back to AABB validation.""" @@ -626,19 +626,19 @@ def _collision_mesh_or_aabb_proxy( def _spheres_penetrate_mesh( self, - source: ObjectBase, + source: PlacementAsset, source_mesh: trimesh.Trimesh, - source_sphere_cache_obj: ObjectBase | None, + source_sphere_cache_obj: PlacementAsset | None, source_applies_yaw: bool, source_uses_pose_yaw: bool, source_pos: torch.Tensor, - target: ObjectBase | CollisionObject, + target: PlacementAsset | CollisionObject, target_mesh: trimesh.Trimesh, target_pos: torch.Tensor, target_uses_pose_yaw: bool, mesh_manager: WarpMeshAndSphereCache, tolerance: float, - orientations: dict[ObjectBase, float] | None, + orientations: dict[PlacementAsset, float] | None, ) -> bool: """True if source's spheres penetrate target's mesh or if BVH returns no-face sentinel. @@ -670,13 +670,13 @@ def _spheres_penetrate_mesh( @staticmethod def _effective_yaw( - obj: ObjectBase | CollisionObject, - orientations: dict[ObjectBase, float] | None, + obj: PlacementAsset | CollisionObject, + orientations: dict[PlacementAsset, float] | None, use_pose_yaw: bool, ) -> float: """Resolve effective Z-yaw from sampled orientations or, when allowed, fixed initial pose.""" if orientations is not None and obj in orientations: - return orientations[cast("ObjectBase", obj)] + return orientations[cast("PlacementAsset", obj)] if not use_pose_yaw: return 0.0 pose = obj.get_initial_pose() @@ -706,11 +706,11 @@ def _pair_aabb_overlaps( @staticmethod def _centers_in_target_frame( centers_local: torch.Tensor, - source_obj: ObjectBase, - target_obj: ObjectBase | CollisionObject, + source_obj: PlacementAsset, + target_obj: PlacementAsset | CollisionObject, source_pos: torch.Tensor, target_pos: torch.Tensor, - orientations: dict[ObjectBase, float] | None, + orientations: dict[PlacementAsset, float] | None, source_applies_yaw: bool = True, source_uses_pose_yaw: bool = True, target_uses_pose_yaw: bool = True, diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py index 17fbc9cddf..2b3e0cfbef 100644 --- a/isaaclab_arena/relations/pooled_object_placer.py +++ b/isaaclab_arena/relations/pooled_object_placer.py @@ -16,8 +16,8 @@ from isaaclab_arena.utils.random import get_rngs if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase from isaaclab_arena.relations.collision_object import CollisionObject + from isaaclab_arena.relations.placement_asset import PlacementAsset @dataclass @@ -70,7 +70,7 @@ class PooledObjectPlacer: def __init__( self, - objects: list[ObjectBase], + objects: list[PlacementAsset], placer_params: ObjectPlacerParams, pool_size: int = 100, num_envs: int | None = None, @@ -314,7 +314,7 @@ def total_remaining(self) -> int: # ------------------------------------------------------------------ @property - def objects(self) -> list[ObjectBase]: + def objects(self) -> list[PlacementAsset]: """All objects (including anchors) participating in relation solving.""" return self._objects diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py index 7da71368e9..f1a3a069b7 100644 --- a/isaaclab_arena/relations/relation_solver.py +++ b/isaaclab_arena/relations/relation_solver.py @@ -22,9 +22,9 @@ from isaaclab_arena.relations.relations import On, Relation, RelationBase, UnaryRelation if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase from isaaclab_arena.relations.collision_object import CollisionObject from isaaclab_arena.relations.mesh_pair_cache import MeshPairCache + from isaaclab_arena.relations.placement_asset import PlacementAsset from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox @@ -54,7 +54,7 @@ def __init__( self._last_position_history: list = [] self._last_loss_per_env: torch.Tensor | None = None self._last_no_overlap_pair_count: int = 0 - self._mesh_orientations: list[dict[ObjectBase, float]] | None = None + self._mesh_orientations: list[dict[PlacementAsset, float]] | None = None self._warned_no_mesh: set[str] = set() self._mesh_manager: WarpMeshAndSphereCache | None = None self._mesh_cache: MeshPairCache | None = None @@ -183,17 +183,17 @@ def _compute_no_overlap_loss( def solve( self, - objects: list[ObjectBase], - initial_positions: list[dict[ObjectBase, tuple[float, float, float]]], - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None, + objects: list[PlacementAsset], + initial_positions: list[dict[PlacementAsset, tuple[float, float, float]]], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox] | None = None, env_bboxes_include_yaw: bool = False, - orientations: list[dict[ObjectBase, float]] | None = None, + orientations: list[dict[PlacementAsset, float]] | None = None, collision_objects: list[CollisionObject] | None = None, - ) -> list[dict[ObjectBase, tuple[float, float, float]]]: + ) -> list[dict[PlacementAsset, tuple[float, float, float]]]: """Solve for optimal positions of all objects. Args: - objects: List of ObjectBase instances. Must include at least one object + objects: Placement assets to solve. Must include at least one asset marked with IsAnchor() which serves as a fixed reference. initial_positions: List of dicts (one per env). Use a single-element list for single-env placement. @@ -353,7 +353,7 @@ def last_position_history(self) -> list: """Position snapshots from the most recent solve() call.""" return self._last_position_history - def debug_losses(self, objects: list[ObjectBase]) -> None: + def debug_losses(self, objects: list[PlacementAsset]) -> None: """Print detailed loss breakdown for all relations using final positions. Call this after solve() to inspect why objects may not be correctly positioned. @@ -378,7 +378,7 @@ def debug_losses(self, objects: list[ObjectBase]) -> None: def _print_relation_debug( - obj: ObjectBase, + obj: PlacementAsset, relation: Relation, child_pos: torch.Tensor, parent_pos: torch.Tensor, @@ -424,7 +424,7 @@ def _print_relation_debug( def _print_unary_relation_debug( - obj: ObjectBase, + obj: PlacementAsset, relation: RelationBase, child_pos: torch.Tensor, loss: torch.Tensor, diff --git a/isaaclab_arena/relations/relation_solver_state.py b/isaaclab_arena/relations/relation_solver_state.py index 377af98ab6..e18ae5a160 100644 --- a/isaaclab_arena/relations/relation_solver_state.py +++ b/isaaclab_arena/relations/relation_solver_state.py @@ -12,8 +12,8 @@ from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase from isaaclab_arena.relations.collision_object import CollisionObject + from isaaclab_arena.relations.placement_asset import PlacementAsset class RelationSolverState: @@ -28,16 +28,16 @@ class RelationSolverState: def __init__( self, - objects: list[ObjectBase], - initial_positions: list[dict[ObjectBase, tuple[float, float, float]]], + objects: list[PlacementAsset], + initial_positions: list[dict[PlacementAsset, tuple[float, float, float]]], device: torch.device | None = None, - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None, + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox] | None = None, collision_objects: list[CollisionObject] | None = None, ): """Initialize optimization state. Args: - objects: List of all ObjectBase instances to track. Must include at least one + objects: Placement assets to track. Must include at least one object marked with IsAnchor() which serves as a fixed reference. initial_positions: List of dicts (one per env). Length 1 = single-env, length > 1 = batched. @@ -52,7 +52,7 @@ def __init__( assert len(anchor_objects) > 0, "No anchor object found in objects list." self._all_objects = objects - self._anchor_objects: set[ObjectBase] = set(anchor_objects) + self._anchor_objects: set[PlacementAsset] = set(anchor_objects) self._optimizable_objects = [obj for obj in objects if obj not in self._anchor_objects] self._collision_objects: list[CollisionObject] = list(collision_objects) if collision_objects else [] assert not (set(self._collision_objects) & set(objects)), ( @@ -61,7 +61,7 @@ def __init__( ) # Build object-to-index mapping - self._obj_to_idx: dict[ObjectBase, int] = {obj: i for i, obj in enumerate(objects)} + self._obj_to_idx: dict[PlacementAsset, int] = {obj: i for i, obj in enumerate(objects)} self._device = device or torch.device("cpu") self._batch_size = len(initial_positions) @@ -110,7 +110,7 @@ def __init__( # Anchors and background collision objects are fixed, so their world bounding boxes are # constant during the solve. Cache them once instead of recomputing every gradient step. - self._fixed_obstacle_world_bboxes: dict[ObjectBase | CollisionObject, AxisAlignedBoundingBox] = { + self._fixed_obstacle_world_bboxes: dict[PlacementAsset | CollisionObject, AxisAlignedBoundingBox] = { obj: obj.get_world_bounding_box().to(self._device) for obj in (*self._anchor_objects, *self._collision_objects) } @@ -134,12 +134,12 @@ def optimizable_positions(self) -> torch.Tensor | None: return self._optimizable_positions @property - def optimizable_objects(self) -> list[ObjectBase]: + def optimizable_objects(self) -> list[PlacementAsset]: """List of optimizable objects (excludes anchors).""" return self._optimizable_objects @property - def anchor_objects(self) -> set[ObjectBase]: + def anchor_objects(self) -> set[PlacementAsset]: """Set of anchor objects (fixed during optimization).""" return self._anchor_objects @@ -148,7 +148,7 @@ def collision_objects(self) -> list[CollisionObject]: """Copy of the collision-only fixed obstacles (constant world pose, no relation constraints).""" return list(self._collision_objects) - def get_position(self, obj: ObjectBase) -> torch.Tensor: + def get_position(self, obj: PlacementAsset) -> torch.Tensor: """Get current position for an object. Args: @@ -169,14 +169,14 @@ def get_position(self, obj: ObjectBase) -> torch.Tensor: opt_idx = self._global_to_opt_idx[idx] return self._optimizable_positions[:, opt_idx, :] - def get_fixed_obstacle_world_bbox(self, obj: ObjectBase | CollisionObject) -> AxisAlignedBoundingBox: + def get_fixed_obstacle_world_bbox(self, obj: PlacementAsset | CollisionObject) -> AxisAlignedBoundingBox: """Return the cached constant world bounding box for an anchor or collision object.""" assert ( obj in self._fixed_obstacle_world_bboxes ), f"'{obj.name}' is not a fixed obstacle (anchor or collision object) tracked by this state." return self._fixed_obstacle_world_bboxes[obj] - def get_bbox(self, obj: ObjectBase) -> AxisAlignedBoundingBox: + def get_bbox(self, obj: PlacementAsset) -> AxisAlignedBoundingBox: """Return the local bounding box for obj, moved to the state's device.""" if self._env_bboxes is not None and obj in self._env_bboxes: return self._env_bboxes[obj].to(self._device) @@ -190,7 +190,7 @@ def get_all_positions_snapshot(self) -> list[tuple[float, float, float]]: """ return [tuple(self.get_position(obj)[0].detach().tolist()) for obj in self._all_objects] - def get_final_positions(self) -> list[dict[ObjectBase, tuple[float, float, float]]]: + def get_final_positions(self) -> list[dict[PlacementAsset, tuple[float, float, float]]]: """Get final positions as a list of dicts, one per env. Returns: diff --git a/isaaclab_arena/relations/relations.py b/isaaclab_arena/relations/relations.py index 287c6ff154..ad13f1953a 100644 --- a/isaaclab_arena/relations/relations.py +++ b/isaaclab_arena/relations/relations.py @@ -15,7 +15,7 @@ from isaaclab_arena.utils.pose import PoseRange # runtime: constructed in to_pose_range_centered_at() if TYPE_CHECKING: - from isaaclab_arena.assets.object_base import ObjectBase + from isaaclab_arena.relations.placement_asset import PlacementAsset RelationT = TypeVar("RelationT", bound="RelationBase") @@ -41,7 +41,7 @@ class RelationBase: in its relations list. """ - def validate_placement_configuration(self, subject: ObjectBase, objects: set[ObjectBase]) -> None: + def validate_placement_configuration(self, subject: PlacementAsset, objects: set[PlacementAsset]) -> None: """Validate this relation for placement among the participating objects.""" @@ -66,7 +66,7 @@ def is_unary() -> bool: """Return whether the relation constrains a single object.""" return False - def __init__(self, parent: ObjectBase, relation_loss_weight: float = 1.0): + def __init__(self, parent: PlacementAsset, relation_loss_weight: float = 1.0): """ Args: parent: The parent asset in the relationship. @@ -86,7 +86,7 @@ class FaceTo(RelationBase): name = "face_to" - def __init__(self, parent: ObjectBase): + def __init__(self, parent: PlacementAsset): """ Args: parent: Target object that defines the facing direction. @@ -98,7 +98,7 @@ def is_unary() -> bool: """Return whether the relation constrains a single object.""" return False - def validate_placement_configuration(self, subject: ObjectBase, objects: set[ObjectBase]) -> None: + def validate_placement_configuration(self, subject: PlacementAsset, objects: set[PlacementAsset]) -> None: """Validate the facing relation for its subject and target.""" face_to_count = sum(isinstance(relation, FaceTo) for relation in subject.get_relations()) assert face_to_count == 1, f"Object '{subject.name}' has more than one FaceTo relation." @@ -140,7 +140,7 @@ class NextTo(Relation): def __init__( self, - parent: ObjectBase, + parent: PlacementAsset, relation_loss_weight: float = 1.0, distance_m: float = 0.05, side: Side | str = Side.POSITIVE_X, @@ -190,7 +190,7 @@ class On(Relation): def __init__( self, - parent: ObjectBase, + parent: PlacementAsset, relation_loss_weight: float = 1.0, clearance_m: float = 0.01, edge_margin_m: float = DEFAULT_ON_EDGE_MARGIN_M, @@ -229,7 +229,7 @@ class NotNextTo(Relation): def __init__( self, - parent: ObjectBase, + parent: PlacementAsset, relation_loss_weight: float = 1.0, side: Side | str = Side.POSITIVE_X, tolerance_m: float = 1e-2, @@ -517,7 +517,7 @@ def __init__( self.relation_loss_weight = relation_loss_weight -def get_anchor_objects(objects: list[ObjectBase]) -> list[ObjectBase]: +def get_anchor_objects(objects: list[PlacementAsset]) -> list[PlacementAsset]: """Get all anchor objects from a list of objects. Anchor objects are marked with IsAnchor() relation and serve as @@ -532,6 +532,6 @@ def get_anchor_objects(objects: list[ObjectBase]) -> list[ObjectBase]: return [obj for obj in objects if any(isinstance(r, IsAnchor) for r in obj.get_relations())] -def get_relation(obj: ObjectBase, relation_type: type[RelationT]) -> RelationT | None: +def get_relation(obj: PlacementAsset, relation_type: type[RelationT]) -> RelationT | None: """Return obj's first relation of the given type, or None if it has none.""" return next((relation for relation in obj.get_relations() if isinstance(relation, relation_type)), None) diff --git a/isaaclab_arena/tasks/lift_object_task.py b/isaaclab_arena/tasks/lift_object_task.py index 4d4ad7f07d..d431d944fb 100644 --- a/isaaclab_arena/tasks/lift_object_task.py +++ b/isaaclab_arena/tasks/lift_object_task.py @@ -192,10 +192,10 @@ def __init__( self.embodiment = embodiment self.observation_cfg = LiftObjectObservationsCfg( - lift_object=self.lift_object, robot_name=self.embodiment.get_embodiment_name_in_scene() + lift_object=self.lift_object, robot_name=self.embodiment.get_scene_name() ) self.commands_cfg = LiftObjectCommandsCfg( - asset_name=self.embodiment.get_embodiment_name_in_scene(), + asset_name=self.embodiment.get_scene_name(), body_name=self.embodiment.get_command_body_name(), lift_object=self.lift_object, target_x_range=self.target_x_range, @@ -205,7 +205,7 @@ def __init__( self.rewards_cfg = LiftObjectRewardCfg( lift_object=self.lift_object, minimum_height_to_lift=self.minimum_height_to_lift, - robot_name=self.embodiment.get_embodiment_name_in_scene(), + robot_name=self.embodiment.get_scene_name(), ee_frame_name=self.embodiment.get_ee_frame_name(self.embodiment.get_arm_mode()), ) @@ -227,7 +227,7 @@ def make_rl_termination_cfg(self): func=lift_object_rl_success, params={ "object_cfg": SceneEntityCfg(self.lift_object.name), - "robot_cfg": SceneEntityCfg(self.embodiment.get_embodiment_name_in_scene()), + "robot_cfg": SceneEntityCfg(self.embodiment.get_scene_name()), "rl_training": self.rl_training_mode, "command_name": "object_pose", "position_tolerance": self.goal_position_tolerance, diff --git a/isaaclab_arena/terms/events.py b/isaaclab_arena/terms/events.py index e0e859667d..2f67971bec 100644 --- a/isaaclab_arena/terms/events.py +++ b/isaaclab_arena/terms/events.py @@ -10,6 +10,7 @@ from isaaclab.managers import SceneEntityCfg from isaaclab_arena.utils.pose import Pose +from isaaclab_arena.utils.scene_pose_writes import write_scene_root_poses_to_sim from isaaclab_arena.utils.velocity import Velocity @@ -63,6 +64,55 @@ def set_object_pose_per_env( ) +def _write_scene_pose(env: ManagerBasedEnv, scene_name: str, pose: Pose, env_ids: torch.Tensor) -> None: + """Write one env-local ``pose`` to a scene entity's root across ``env_ids`` (env origins added).""" + asset = env.scene[scene_name] + num_envs = len(env_ids) + pose_t_xyz_q_xyzw = pose.to_tensor(device=env.device).repeat(num_envs, 1) + pose_t_xyz_q_xyzw[:, :3] += env.scene.env_origins[env_ids] + write_scene_root_poses_to_sim(asset, scene_name, pose_t_xyz_q_xyzw, env_ids, env.device) + + +def reset_placement_asset_pose( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + scene_writes: list[tuple[str, Pose]], +) -> None: + """Restore a placement asset to fixed env-local poses on reset. + + Each ``(scene entity name, pose)`` in ``scene_writes`` is written to every resetting env, + letting a compound asset place several prims (e.g. a robot and its stand) together. + """ + if env_ids is None: + return + for scene_name, pose in scene_writes: + _write_scene_pose(env, scene_name, pose, env_ids) + + +def reset_placement_asset_pose_per_env( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + write_pose_list: list[list[tuple[str, Pose]]], +) -> None: + """Restore a placement asset to a distinct per-env pose on reset. + + ``write_pose_list[env]`` holds that env's ``(scene entity name, pose)`` writes, so different + environments can hold different solved layouts for the same asset (and its auxiliary prims). + """ + if env_ids is None: + return + assert env_ids.ndim == 1, "env_ids must be a 1-D tensor of environment indices" + num_scene_envs = env.scene.env_origins.shape[0] + assert len(write_pose_list) == num_scene_envs, ( + f"per-env pose writes cover {len(write_pose_list)} envs, but the scene has {num_scene_envs}; " + "write_pose_list is indexed by absolute env id and must span every environment." + ) + for cur_env in env_ids.tolist(): + single_env = torch.tensor([cur_env], device=env.device) + for scene_name, pose in write_pose_list[cur_env]: + _write_scene_pose(env, scene_name, pose, single_env) + + def reset_all_articulation_joints(env: ManagerBasedEnv, env_ids: torch.Tensor): """Reset the articulation joints to the initial state.""" for articulation_asset in env.scene.articulations.values(): diff --git a/isaaclab_arena/tests/dummy_embodiment.py b/isaaclab_arena/tests/dummy_embodiment.py new file mode 100644 index 0000000000..4f83633d6e --- /dev/null +++ b/isaaclab_arena/tests/dummy_embodiment.py @@ -0,0 +1,81 @@ +# Copyright (c) 2026, 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 + +"""Lightweight placement embodiment for solver tests.""" + +from __future__ import annotations + +import trimesh +from typing import TYPE_CHECKING + +from isaaclab_arena.relations.placement_asset import PlacementAsset +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.pose import Pose, PosePerEnv + +if TYPE_CHECKING: + from isaaclab.managers import EventTermCfg + + +class DummyEmbodiment(PlacementAsset): + """Embodiment geometry without simulator dependencies. + + Mirrors the real embodiment pose lifecycle: setting an initial pose builds a + root-pose reset event, so the solver's ``has_pose_reset_event`` invariant holds. + """ + + def __init__( + self, + name: str, + bounding_box: AxisAlignedBoundingBox, + initial_pose: Pose | None = None, + collision_mesh: trimesh.Trimesh | None = None, + scene_name: str | None = None, + ) -> None: + super().__init__(name=name, tags=["embodiment"]) + self.initial_pose = initial_pose + self.bounding_box = bounding_box + self.collision_mesh = collision_mesh + self.scene_name = name if scene_name is None else scene_name + self.pose_event_cfg: EventTermCfg | None = None + + def set_initial_pose(self, pose: Pose | PosePerEnv) -> None: + """Store the requested pose(s) and rebuild the root-pose reset event.""" + self.initial_pose = pose + self.pose_event_cfg = self._init_pose_event_cfg() + + def _init_pose_event_cfg(self) -> EventTermCfg | None: + from isaaclab.managers import EventTermCfg + + from isaaclab_arena.terms.events import reset_placement_asset_pose, reset_placement_asset_pose_per_env + + if self.initial_pose is None: + return None + if isinstance(self.initial_pose, PosePerEnv): + return EventTermCfg( + func=reset_placement_asset_pose_per_env, + mode="reset", + params={"write_pose_list": [self.layout_pose_to_scene_writes(p) for p in self.initial_pose.poses]}, + ) + return EventTermCfg( + func=reset_placement_asset_pose, + mode="reset", + params={"scene_writes": self.layout_pose_to_scene_writes(self.initial_pose)}, + ) + + def has_pose_reset_event(self) -> bool: + """Return whether a pose was set and therefore a reset event exists.""" + return self.pose_event_cfg is not None + + def get_bounding_box(self) -> AxisAlignedBoundingBox: + """Return root-relative bounds.""" + return self.bounding_box + + def get_collision_mesh(self) -> trimesh.Trimesh | None: + """Return the configured collision mesh.""" + return self.collision_mesh + + def get_scene_name(self) -> str: + """Return the configured scene key.""" + return self.scene_name diff --git a/isaaclab_arena/tests/dummy_object.py b/isaaclab_arena/tests/dummy_object.py new file mode 100644 index 0000000000..0977931360 --- /dev/null +++ b/isaaclab_arena/tests/dummy_object.py @@ -0,0 +1,84 @@ +# Copyright (c) 2025-2026, 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 __future__ import annotations + +import torch +import trimesh +from typing import TYPE_CHECKING + +from isaaclab_arena.relations.placement_asset import PlacementAsset +from isaaclab_arena.relations.relations import RelationBase +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.pose import Pose, PosePerEnv + +if TYPE_CHECKING: + from isaaclab.managers import EventTermCfg + + +class DummyObject(PlacementAsset): + """Dummy object for testing purposes without Isaac Sim dependencies. + + Mirrors the real object pose lifecycle: setting an initial pose builds a + root-pose reset event, so the solver's ``has_pose_reset_event`` invariant holds. + """ + + def __init__( + self, + name: str, + bounding_box: AxisAlignedBoundingBox, + initial_pose: Pose | None = None, + relations: list[RelationBase] | None = None, + collision_mesh: trimesh.Trimesh | None = None, + **kwargs, + ): + super().__init__(name=name) + self.initial_pose = initial_pose + self.bounding_box = bounding_box + assert self.bounding_box is not None + self.relations = list(relations or []) + self._collision_mesh = collision_mesh + self.pose_event_cfg: EventTermCfg | None = None + + def set_initial_pose(self, pose: Pose | PosePerEnv) -> None: + """Store the requested pose(s) and rebuild the root-pose reset event.""" + self.initial_pose = pose + self.pose_event_cfg = self._init_pose_event_cfg() + + def _init_pose_event_cfg(self) -> EventTermCfg | None: + from isaaclab.managers import EventTermCfg + + from isaaclab_arena.terms.events import reset_placement_asset_pose, reset_placement_asset_pose_per_env + + if self.initial_pose is None: + return None + if isinstance(self.initial_pose, PosePerEnv): + return EventTermCfg( + func=reset_placement_asset_pose_per_env, + mode="reset", + params={"write_pose_list": [self.layout_pose_to_scene_writes(p) for p in self.initial_pose.poses]}, + ) + return EventTermCfg( + func=reset_placement_asset_pose, + mode="reset", + params={"scene_writes": self.layout_pose_to_scene_writes(self.initial_pose)}, + ) + + def has_pose_reset_event(self) -> bool: + """Return whether a pose was set and therefore a reset event exists.""" + return self.pose_event_cfg is not None + + def get_bounding_box(self) -> AxisAlignedBoundingBox: + """Get local bounding box (relative to object origin).""" + return self.bounding_box + + def get_corners_aabb(self, pos: torch.Tensor) -> torch.Tensor: + return self.bounding_box.get_corners_at(pos) + + def is_initial_pose_set(self) -> bool: + return self.initial_pose is not None + + def get_collision_mesh(self) -> trimesh.Trimesh | None: + """Return the collision mesh, or None to fall back to AABB.""" + return self._collision_mesh diff --git a/isaaclab_arena/tests/test_droid_stand_height.py b/isaaclab_arena/tests/test_droid_stand_height.py index 300d40c2d9..d1996bd41d 100644 --- a/isaaclab_arena/tests/test_droid_stand_height.py +++ b/isaaclab_arena/tests/test_droid_stand_height.py @@ -3,11 +3,19 @@ # # SPDX-License-Identifier: Apache-2.0 +import torch import traceback +from pathlib import Path from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function _CUSTOM_STAND_HEIGHT_M = 2.0 +_KITCHEN_STAND_HEIGHT_M = 0.8 +_KITCHEN_YAML = ( + Path(__file__).resolve().parents[2] / "isaaclab_arena_environments" / "droid_pick_and_place_lightwheel_kitchen.yaml" +) +_Z_MATCH_EPS = 1e-3 +_QUAT_MATCH_EPS = 1e-3 def _test_droid_stand_height(simulation_app) -> bool: @@ -55,10 +63,22 @@ def _test_droid_stand_height(simulation_app) -> bool: assert abs(posed_emb.initial_pose.position_xyz[2] - (0.5 + expected_offset)) < 1e-6 scene_cfg = posed_emb.get_scene_cfg() assert abs(scene_cfg.robot.init_state.pos[2] - (0.5 + expected_offset)) < 1e-6 + assert abs(scene_cfg.stand.init_state.pos[0] - (0.3 - 0.05)) < 1e-6 assert abs(scene_cfg.stand.init_state.pos[2] - (0.5 + expected_offset)) < 1e-6 # initial_pose is the single source of truth: it matches the spawned robot base exactly. assert tuple(posed_emb.initial_pose.position_xyz) == tuple(scene_cfg.robot.init_state.pos) + layout_pose = Pose(position_xyz=(0.3, 0.0, 0.5), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) + scene_writes = posed_emb.layout_pose_to_scene_writes(layout_pose) + assert len(scene_writes) == 2 + robot_scene_name, robot_write_pose = scene_writes[0] + stand_scene_name, stand_write_pose = scene_writes[1] + assert robot_scene_name == "robot" + assert stand_scene_name == "stand" + assert abs(robot_write_pose.position_xyz[2] - (0.5 + expected_offset)) < 1e-6 + assert abs(stand_write_pose.position_xyz[0] - (0.3 - 0.05)) < 1e-6 + assert abs(stand_write_pose.position_xyz[2] - robot_write_pose.position_xyz[2]) < 1e-6 + # The YAML-spec path instantiates the embodiment via asset_class(**params); a scalar # stand_height_m from YAML applies just the same. registry_emb = AssetRegistry().get_asset_by_name("droid_abs_joint_pos")(stand_height_m=_CUSTOM_STAND_HEIGHT_M) @@ -73,11 +93,113 @@ def _test_droid_stand_height(simulation_app) -> bool: return True +def _test_droid_placement_bbox(simulation_app) -> bool: + """Check Droid placement bounds come from the stand footprint, not the robot mesh.""" + + from isaaclab_arena.embodiments.droid.droid import ( + _STAND_ROOT_OFFSET_IN_ROBOT_FRAME, + DroidAbsoluteJointPositionEmbodiment, + ) + from isaaclab_arena.utils.pose import Pose, translate_by_xyz_offset + from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd + + try: + embodiment = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_KITCHEN_STAND_HEIGHT_M) + bbox = embodiment.get_bounding_box() + + stand = embodiment.scene_config.stand + stand_bbox = compute_local_bounding_box_from_usd(stand.spawn.usd_path, tuple(stand.spawn.scale)) + stand_offset = translate_by_xyz_offset(_STAND_ROOT_OFFSET_IN_ROBOT_FRAME, embodiment._robot_base_offset) + expected_bbox = stand_bbox.translated(stand_offset) + assert torch.allclose(bbox.min_point, expected_bbox.min_point) + assert torch.allclose(bbox.max_point, expected_bbox.max_point) + + robot_bbox = compute_local_bounding_box_from_usd( + embodiment.scene_config.robot.spawn.usd_path, + tuple(embodiment.scene_config.robot.spawn.scale or (1.0, 1.0, 1.0)), + ) + assert bbox.size[0, 0].item() > robot_bbox.size[0, 0].item() + assert bbox.size[0, 1].item() > robot_bbox.size[0, 1].item() + assert bbox.min_point[0, 2].item() <= 0.0 + + # Relation solver passes unlifted poses; set_initial_pose applies the stand-height offset. + floor_top_z = 0.0 + solver_z = floor_top_z - bbox.min_point[0, 2].item() + embodiment.set_initial_pose(Pose(position_xyz=(0.0, 0.0, solver_z), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + spawn_z = embodiment.initial_pose.position_xyz[2] + natural_stand_bottom_z = spawn_z + (expected_bbox.min_point[0, 2].item() - stand_offset[2]) + assert abs(natural_stand_bottom_z - floor_top_z) < _Z_MATCH_EPS + except Exception as exc: + print(f"Error: {exc}") + traceback.print_exc() + return False + + return True + + +def _test_droid_kitchen_reset_alignment(simulation_app) -> bool: + """Verify Droid robot and stand stay z-aligned after relation placement reset.""" + + import warp as wp + + from isaaclab_arena.environment_spec.arena_env_graph_spec import ArenaEnvGraphSpec + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.arena_env_builder_cfg import ArenaEnvBuilderCfg + + try: + spec = ArenaEnvGraphSpec.from_yaml(_KITCHEN_YAML) + params = dict(spec.embodiment.params) + params["stand_height_m"] = _KITCHEN_STAND_HEIGHT_M + spec.embodiment.params = params + + builder = ArenaEnvBuilder(spec.to_arena_env(), ArenaEnvBuilderCfg(num_envs=1)) + env = builder.make_registered() + env.reset() + + robot_pose = wp.to_torch(env.unwrapped.scene["robot"].data.root_link_pose_w)[0] + robot_z = robot_pose[2].item() + robot_quat = robot_pose[3:7] + + stand_positions, stand_orientations = env.unwrapped.scene["stand"].get_world_poses() + stand_z = stand_positions[0, 2].item() + stand_quat = stand_orientations[0] + + assert abs(robot_z - stand_z) < _Z_MATCH_EPS, f"robot z {robot_z} != stand z {stand_z} after reset" + assert torch.allclose( + robot_quat.cpu(), + stand_quat.cpu(), + atol=_QUAT_MATCH_EPS, + rtol=0.0, + ), f"robot quat {robot_quat.tolist()} != stand quat {stand_quat.tolist()} after reset" + + env.close() + except Exception as exc: + print(f"Error: {exc}") + traceback.print_exc() + return False + + return True + + def test_droid_stand_height(): """Pytest entry point for the Droid stand-height configuration test.""" result = run_simulation_app_function(_test_droid_stand_height, headless=True) assert result, f"Test {test_droid_stand_height.__name__} failed" +def test_droid_placement_bbox(): + """Pytest entry point for the Droid stand placement-bbox test.""" + result = run_simulation_app_function(_test_droid_placement_bbox, headless=True) + assert result, f"Test {test_droid_placement_bbox.__name__} failed" + + +def test_droid_kitchen_reset_alignment(): + """Pytest entry point for Droid kitchen reset alignment.""" + result = run_simulation_app_function(_test_droid_kitchen_reset_alignment, headless=True) + assert result, f"Test {test_droid_kitchen_reset_alignment.__name__} failed" + + if __name__ == "__main__": test_droid_stand_height() + test_droid_placement_bbox() + test_droid_kitchen_reset_alignment() diff --git a/isaaclab_arena/tests/test_embodiment_collision_mesh.py b/isaaclab_arena/tests/test_embodiment_collision_mesh.py new file mode 100644 index 0000000000..9b9487d2e6 --- /dev/null +++ b/isaaclab_arena/tests/test_embodiment_collision_mesh.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026, 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 traceback + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + + +def _test_embodiment_provides_robot_collision_mesh(simulation_app) -> bool: + """Check the embodiment exposes its robot mesh so MESH mode does not fall back to the bbox proxy.""" + + from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment + + try: + emb = DroidAbsoluteJointPositionEmbodiment() + + mesh = emb.get_collision_mesh() + assert mesh is not None, "embodiment must expose a collision mesh; None forces the loose bbox fallback" + assert len(mesh.vertices) > 0 + + # The default prim scopes extraction to the arm. The Droid USD also bakes in a 50 m ground + # plane and stray props; leaking those would blow the mesh up to scene scale. + extents = mesh.extents + assert all(e < 2.0 for e in extents), f"mesh leaked non-robot geometry: extents {extents}" + + # The mesh and the bounding box describe the same body, so their extents track each other. + bbox = emb.get_bounding_box() + bbox_size = (bbox.max_point - bbox.min_point)[0].tolist() + for mesh_extent, box_extent in zip(extents, bbox_size): + assert abs(mesh_extent - box_extent) < 0.2, f"mesh extents {extents} disagree with bbox {bbox_size}" + + # Extraction opens the USD, so the result is cached rather than recomputed per solve. + assert emb.get_collision_mesh() is mesh + + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + + return True + + +def test_embodiment_provides_robot_collision_mesh(): + """Pytest entry point for the embodiment collision-mesh test.""" + result = run_simulation_app_function(_test_embodiment_provides_robot_collision_mesh, headless=True) + assert result, f"Test {test_embodiment_provides_robot_collision_mesh.__name__} failed" + + +if __name__ == "__main__": + test_embodiment_provides_robot_collision_mesh() diff --git a/isaaclab_arena/tests/test_face_to.py b/isaaclab_arena/tests/test_face_to.py index c880bf5433..f5671a6454 100644 --- a/isaaclab_arena/tests/test_face_to.py +++ b/isaaclab_arena/tests/test_face_to.py @@ -13,7 +13,6 @@ import pytest from pydantic import ValidationError -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.assets.registries import ObjectRelationLibraryRegistry from isaaclab_arena.environment_spec.arena_env_graph_types import SpatialRelationSpec from isaaclab_arena.relations.collision_mode import CollisionMode @@ -24,6 +23,7 @@ from isaaclab_arena.relations.relation_solver import RelationSolver from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import AtPosition, FaceTo, IsAnchor, RandomAroundSolution, RotateAroundSolution +from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange from isaaclab_arena.utils.yaw import wrap_angle_to_pi, yaw_from_quat_xyzw, yaw_toward_positions diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py index 5a36dfe931..77478a43b1 100644 --- a/isaaclab_arena/tests/test_heterogeneous_placement.py +++ b/isaaclab_arena/tests/test_heterogeneous_placement.py @@ -13,7 +13,6 @@ import pytest -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.bounding_box_helpers import build_per_env_bounding_boxes, get_bounding_box_per_env from isaaclab_arena.relations.object_placer import ObjectPlacer from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams @@ -23,6 +22,7 @@ from isaaclab_arena.relations.relation_solver import RelationSolver from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On +from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose diff --git a/isaaclab_arena/tests/test_mesh_collision.py b/isaaclab_arena/tests/test_mesh_collision.py index 75bda87c6f..11b1e4cdde 100644 --- a/isaaclab_arena/tests/test_mesh_collision.py +++ b/isaaclab_arena/tests/test_mesh_collision.py @@ -14,12 +14,12 @@ import pytest -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.relation_loss_strategies import NoCollisionLossStrategy from isaaclab_arena.relations.relation_solver import RelationSolver from isaaclab_arena.relations.relation_solver_params import CollisionMode, RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On from isaaclab_arena.relations.warp_mesh_manager import WarpMeshAndSphereCache, greedy_sphere_decomposition +from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose diff --git a/isaaclab_arena/tests/test_no_collision_loss.py b/isaaclab_arena/tests/test_no_collision_loss.py index 89fb06d08e..3b17b5fb33 100644 --- a/isaaclab_arena/tests/test_no_collision_loss.py +++ b/isaaclab_arena/tests/test_no_collision_loss.py @@ -8,13 +8,13 @@ import math import torch -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.loss_primitives import interval_overlap_axis_loss from isaaclab_arena.relations.relation_loss_strategies import NoCollisionLossStrategy from isaaclab_arena.relations.relation_solver import RelationSolver from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relation_solver_state import RelationSolverState from isaaclab_arena.relations.relations import IsAnchor, On +from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose diff --git a/isaaclab_arena/tests/test_object_placer_init.py b/isaaclab_arena/tests/test_object_placer_init.py index 3510a03c5a..97104e3675 100644 --- a/isaaclab_arena/tests/test_object_placer_init.py +++ b/isaaclab_arena/tests/test_object_placer_init.py @@ -5,11 +5,11 @@ """Tests for On-relation-guided initialization in ObjectPlacer.""" -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.object_placer import ObjectPlacer from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, NextTo, On, Side +from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose diff --git a/isaaclab_arena/tests/test_object_placer_reproducibility.py b/isaaclab_arena/tests/test_object_placer_reproducibility.py index c01cca5011..08eb807283 100644 --- a/isaaclab_arena/tests/test_object_placer_reproducibility.py +++ b/isaaclab_arena/tests/test_object_placer_reproducibility.py @@ -9,7 +9,6 @@ import pytest -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.object_placer import ObjectPlacer from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.placement_result import PlacementResult @@ -17,6 +16,7 @@ from isaaclab_arena.relations.relation_solver import RelationSolver from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, NextTo, On, RotateAroundSolution, Side +from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, get_random_pose_within_bounding_box from isaaclab_arena.utils.pose import Pose, PosePerEnv from isaaclab_arena.utils.yaw import rotate_quat_by_yaw, wrap_angle_to_pi diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index 496b2e4fb7..7176a3ee3d 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -25,8 +25,8 @@ def _checklist(passed: bool): def _create_test_objects(): """Create a desk (anchor) with two boxes (On + NextTo).""" - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.relations import IsAnchor, NextTo, On, Side + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -54,8 +54,6 @@ def _create_test_objects(): def test_successive_placements_without_seed_produce_different_layouts(): - """Two place() calls with placement_seed=None should produce different positions.""" - from isaaclab_arena.relations.object_placer import ObjectPlacer from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -84,8 +82,6 @@ def test_successive_placements_without_seed_produce_different_layouts(): def test_placement_without_seed_multi_env_gives_different_layouts(): - """Multi-env placement with seed=None should give distinct per-env layouts.""" - from isaaclab_arena.relations.object_placer import ObjectPlacer from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -109,8 +105,6 @@ def test_placement_without_seed_multi_env_gives_different_layouts(): def test_successive_seeded_placements_produce_same_layout(): - """Two place() calls with the same seed should produce identical positions.""" - from isaaclab_arena.relations.object_placer import ObjectPlacer from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -155,6 +149,11 @@ def scene_getitem(self, name: str) -> MagicMock: return env +def _configure_asset_base_mock(asset_mock: MagicMock) -> None: + """Make a mock scene asset behave like AssetBase (``set_world_poses`` only).""" + asset_mock.write_root_pose_to_sim = None + + def _solve_and_place_with_pool(env, env_ids, objects, pool): """Call the reset event with the same runtime params EventTermCfg stores.""" from isaaclab_arena.relations.placement_events import solve_and_place_objects @@ -162,14 +161,12 @@ def _solve_and_place_with_pool(env, env_ids, objects, pool): return solve_and_place_objects( env, env_ids, - objects=objects, + assets=objects, placement_pool=pool, ) def test_solve_and_place_objects_writes_poses_to_sim(): - """solve_and_place_objects should call write_root_pose_to_sim for non-anchor objects.""" - from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -199,11 +196,17 @@ def test_solve_and_place_objects_writes_poses_to_sim(): def test_solve_and_place_objects_uses_runtime_pool(): - """Reset params should use the runtime placement pool directly.""" from isaaclab_arena.relations.placement_events import solve_and_place_objects from isaaclab_arena.relations.placement_result import PlacementResult + from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox - desk, box1, _ = _create_test_objects() + desk, _, _ = _create_test_objects() + robot = DummyEmbodiment( + name="droid", + bounding_box=AxisAlignedBoundingBox(min_point=(-0.2, -0.2, 0.0), max_point=(0.2, 0.2, 1.0)), + scene_name="robot", + ) env = _make_mock_env(num_envs=1) class Pool: @@ -214,7 +217,7 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: return { 0: PlacementResult( validation_results=_checklist(True), - positions={box1: (0.2, 0.3, 0.4)}, + positions={robot: (0.2, 0.3, 0.4)}, final_loss=0.0, attempts=1, ) @@ -223,16 +226,140 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: solve_and_place_objects( env, torch.tensor([0]), - objects=[desk, box1], + assets=[desk, robot], placement_pool=Pool(), ) assert "desk" not in env._assets - env._assets["box1"].write_root_pose_to_sim.assert_called_once() + assert "droid" not in env._assets + env._assets["robot"].write_root_pose_to_sim.assert_called_once() + + +def _identity_pose(position_xyz): + from isaaclab_arena.utils.pose import Pose + + return Pose(position_xyz=position_xyz, rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) + + +def test_reset_placement_asset_pose_writes_compound_prims_with_env_origins(): + """A single fixed layout writes every scene entity to all resetting envs, origin-shifted, zero velocity.""" + from isaaclab_arena.terms.events import reset_placement_asset_pose + + env = _make_mock_env(num_envs=2) + env.scene.env_origins = torch.tensor([[10.0, 0.0, 0.0], [0.0, 20.0, 0.0]]) + scene_writes = [("robot", _identity_pose((0.1, 0.2, 0.5))), ("stand", _identity_pose((0.1, 0.2, 0.0)))] + + reset_placement_asset_pose(env, torch.tensor([0, 1]), scene_writes=scene_writes) + + for name, base in (("robot", (0.1, 0.2, 0.5)), ("stand", (0.1, 0.2, 0.0))): + pose = env._assets[name].write_root_pose_to_sim.call_args.args[0] + assert pose.shape == (2, 7) + assert torch.allclose(pose[0, :3], torch.tensor([base[0] + 10.0, base[1], base[2]])) + assert torch.allclose(pose[1, :3], torch.tensor([base[0], base[1] + 20.0, base[2]])) + velocity = env._assets[name].write_root_velocity_to_sim.call_args.args[0] + assert torch.count_nonzero(velocity) == 0 + + +def test_solve_and_place_objects_writes_droid_robot_and_stand(): + """write_layout_to_sim writes every entry from layout_pose_to_scene_writes (robot + stand).""" + from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment + from isaaclab_arena.relations.placement_events import solve_and_place_objects + from isaaclab_arena.relations.placement_result import PlacementResult + + desk, _, _ = _create_test_objects() + droid = DroidAbsoluteJointPositionEmbodiment(stand_height_m=0.8) + expected_offset = 0.8 - 1.35 + env = _make_mock_env(num_envs=1) + _configure_asset_base_mock(env.scene["stand"]) + + class Pool: + num_envs = 1 + + def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: + return { + 0: PlacementResult( + validation_results=_checklist(True), + positions={droid: (0.2, 0.3, 1.36)}, + final_loss=0.0, + attempts=1, + ) + } + + solve_and_place_objects( + env, + torch.tensor([0]), + assets=[desk, droid], + placement_pool=Pool(), + ) + + robot_pose = env._assets["robot"].write_root_pose_to_sim.call_args.args[0] + stand_call = env._assets["stand"].set_world_poses.call_args + stand_positions = stand_call.kwargs["positions"] + expected_robot_z = 1.36 + expected_offset + assert abs(robot_pose[0, 2].item() - expected_robot_z) < 1e-5 + assert abs(stand_positions[0, 2].item() - expected_robot_z) < 1e-5 + assert abs(stand_positions[0, 0].item() - (0.2 - 0.05)) < 1e-5 + + +def test_write_layout_to_sim_rejects_missing_non_anchor_assets(): + from isaaclab_arena.relations.placement_events import write_layout_to_sim + from isaaclab_arena.relations.placement_result import PlacementResult + + _, box1, _ = _create_test_objects() + env = _make_mock_env(num_envs=1) + layout = PlacementResult( + validation_results=_checklist(True), + positions={}, + final_loss=0.0, + attempts=1, + ) + + with pytest.raises(AssertionError, match="missing non-anchor assets.*box1"): + write_layout_to_sim(env, 0, layout, anchor_assets=set(), base_rotations={box1: (0.0, 0.0, 0.0, 1.0)}) + + +def test_reset_placement_asset_pose_per_env_indexes_by_absolute_env(): + """A partial reset must apply write_pose_list[env_id], not the first layout.""" + from isaaclab_arena.terms.events import reset_placement_asset_pose_per_env + + env = _make_mock_env(num_envs=3) + write_pose_list = [[("robot", _identity_pose((float(i), 0.0, 0.0)))] for i in range(3)] + + reset_placement_asset_pose_per_env(env, torch.tensor([2]), write_pose_list=write_pose_list) + + robot = env._assets["robot"] + robot.write_root_pose_to_sim.assert_called_once() + assert torch.equal(robot.write_root_pose_to_sim.call_args.kwargs["env_ids"], torch.tensor([2])) + assert torch.allclose(robot.write_root_pose_to_sim.call_args.args[0][0, :3], torch.tensor([2.0, 0.0, 0.0])) + + +def test_reset_placement_asset_pose_per_env_writes_each_compound_prim_per_env(): + from isaaclab_arena.terms.events import reset_placement_asset_pose_per_env + + env = _make_mock_env(num_envs=2) + _configure_asset_base_mock(env.scene["stand"]) + write_pose_list = [ + [("robot", _identity_pose((x, 0.0, 0.5))), ("stand", _identity_pose((x, 0.0, 0.0)))] for x in (0.0, 1.0) + ] + + reset_placement_asset_pose_per_env(env, torch.tensor([0, 1]), write_pose_list=write_pose_list) + + assert env._assets["robot"].write_root_pose_to_sim.call_count == 2 + assert env._assets["stand"].set_world_poses.call_count == 2 + + +def test_reset_placement_asset_pose_per_env_requires_full_env_coverage(): + """Guarding the length turns an out-of-range env index into a clear error, not an IndexError.""" + from isaaclab_arena.terms.events import reset_placement_asset_pose_per_env + + env = _make_mock_env(num_envs=3) + short_list = [[("robot", _identity_pose((0.0, 0.0, 0.0)))]] + + with pytest.raises(AssertionError, match="per-env pose writes"): + reset_placement_asset_pose_per_env(env, torch.tensor([2]), write_pose_list=short_list) def test_get_placement_pool_returns_runtime_pool(): - """Pool validation can retrieve the runtime placement pool from the reset event.""" from isaaclab_arena.relations.placement_events import get_placement_pool class Pool: @@ -282,8 +409,6 @@ def test_solve_and_place_objects_applies_random_yaw(): def test_solve_and_place_objects_skips_empty_env_ids(): - """solve_and_place_objects should return immediately for an empty env_ids tensor.""" - from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -301,8 +426,6 @@ def test_solve_and_place_objects_skips_empty_env_ids(): def test_solve_and_place_objects_skips_none_env_ids(): - """solve_and_place_objects should return immediately when env_ids is None.""" - from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -320,8 +443,6 @@ def test_solve_and_place_objects_skips_none_env_ids(): def test_solve_and_place_objects_handles_multiple_env_ids(): - """solve_and_place_objects should write poses for each resetting environment.""" - from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -405,14 +526,13 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: assert "Writing best-loss fallback placement for env 0; failed checks: ['valid']." in captured.out -def test_solve_and_place_objects_partial_reset_env_indexed_uses_absolute_env_result(): - """Env-indexed partial resets should write the result for each absolute env id.""" - +def test_solve_and_place_objects_partial_reset_applies_absolute_env_origin(): from isaaclab_arena.relations.placement_result import PlacementResult desk, box1, box2 = _create_test_objects() objects = [desk, box1, box2] env = _make_mock_env(num_envs=4) + env.scene.env_origins[2] = torch.tensor([10.0, 0.0, 0.0]) class EnvIndexedPool: num_envs = 4 @@ -443,8 +563,8 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: box2_pose = env._assets[box2.name].write_root_pose_to_sim.call_args[0][0] box1_env_id = env._assets[box1.name].write_root_pose_to_sim.call_args.kwargs["env_ids"] box2_env_id = env._assets[box2.name].write_root_pose_to_sim.call_args.kwargs["env_ids"] - assert box1_pose[0, 0].item() == 2.0 - assert box2_pose[0, 0].item() == 2.0 + assert box1_pose[0, 0].item() == 12.0 + assert box2_pose[0, 0].item() == 12.0 assert box1_env_id.tolist() == [2] assert box2_env_id.tolist() == [2] assert pool.requested_env_ids == [2] @@ -465,8 +585,6 @@ class MismatchedEnvIndexedPool: def test_pooled_placer_sample_without_replacement_returns_different_layouts(): - """sample_without_replacement() should return layouts (likely different across draws).""" - from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -488,8 +606,6 @@ def test_pooled_placer_sample_without_replacement_returns_different_layouts(): def test_pooled_object_placer_sample_with_replacement_does_not_consume(): - """sample_with_replacement() should return layouts without consuming them.""" - from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -507,8 +623,6 @@ def test_pooled_object_placer_sample_with_replacement_does_not_consume(): def test_pooled_object_placer_sample_without_replacement_triggers_refill(): - """Exhausting the pool and requesting more should trigger a refill.""" - from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams @@ -581,6 +695,10 @@ def get_relations(self): def set_initial_pose(self, pose): raise AssertionError("resolve_on_reset init seeding must not register per-object reset events") + def set_spawn_pose(self, pose): + self.object_cfg.init_state.pos = pose.position_xyz + self.object_cfg.init_state.rot = pose.rotation_xyzw + class EnvIndexedPool: num_envs = 3 sample_count = None @@ -595,7 +713,7 @@ def sample_with_replacement(self, count: int): final_loss=0.0, attempts=1, ) - for env_id in range(self.num_envs) + for env_id in range(count) ] anchor = MinimalObject("desk") @@ -603,9 +721,9 @@ def sample_with_replacement(self, count: int): pool = EnvIndexedPool() _apply_dynamic_spawn_pose( - objects=[anchor, box], + assets=[anchor, box], placement_pool=pool, - anchor_objects_set={anchor}, + anchor_assets={anchor}, ) assert pool.sample_count == 1 @@ -616,10 +734,10 @@ def sample_with_replacement(self, count: int): def test_env_indexed_static_poses_apply_per_env_positions(): """Static initial poses should apply per-env positions from env-indexed layouts.""" - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.environments.relation_solver_interface import _apply_static_initial_poses from isaaclab_arena.relations.placement_result import PlacementResult from isaaclab_arena.relations.relations import IsAnchor, On + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose, PosePerEnv @@ -653,9 +771,9 @@ def sample_with_replacement(self, count: int): ] _apply_static_initial_poses( - objects=[desk, box], + assets=[desk, box], placement_pool=PerEnvPool(), - anchor_objects_set={desk}, + anchor_assets={desk}, num_envs=num_envs, ) @@ -669,11 +787,11 @@ def sample_with_replacement(self, count: int): def test_pooled_placer_falls_back_when_no_valid_layouts(capsys): """PooledObjectPlacer should keep best-loss fallback layouts when validation rejects all candidates.""" - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -711,11 +829,11 @@ def test_pooled_placer_falls_back_when_no_valid_layouts(capsys): def test_pooled_placer_only_falls_back_on_final_batch(capsys): """Fallbacks should only be accepted on the last configured solve batch.""" - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -747,11 +865,11 @@ def test_pooled_placer_only_falls_back_on_final_batch(capsys): def test_pooled_placer_can_reject_best_loss_fallbacks(): """PooledObjectPlacer should fail loudly when fallback layouts are disabled.""" - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose diff --git a/isaaclab_arena/tests/test_position_limits.py b/isaaclab_arena/tests/test_position_limits.py index a8897398f1..bc398a3c3d 100644 --- a/isaaclab_arena/tests/test_position_limits.py +++ b/isaaclab_arena/tests/test_position_limits.py @@ -156,10 +156,10 @@ def test_position_limits_unconstrained_axes_ignored(): # ============================================================================= -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.relation_solver import RelationSolver from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On +from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.pose import Pose diff --git a/isaaclab_arena/tests/test_relation_loss_strategies.py b/isaaclab_arena/tests/test_relation_loss_strategies.py index 20e2f392cb..b2558cc628 100644 --- a/isaaclab_arena/tests/test_relation_loss_strategies.py +++ b/isaaclab_arena/tests/test_relation_loss_strategies.py @@ -9,11 +9,11 @@ import pytest -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.relation_loss_strategies import NextToLossStrategy, NotNextToLossStrategy, OnLossStrategy from isaaclab_arena.relations.relation_solver import RelationSolver from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, NextTo, NotNextTo, On, Side +from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose diff --git a/isaaclab_arena/tests/test_relation_placement_from_yaml.py b/isaaclab_arena/tests/test_relation_placement_from_yaml.py new file mode 100644 index 0000000000..2fbbc3e8d1 --- /dev/null +++ b/isaaclab_arena/tests/test_relation_placement_from_yaml.py @@ -0,0 +1,79 @@ +# Copyright (c) 2026, 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 + +"""Integration test: build an env from a spec YAML and check the solved layout honors its relations.""" + +import traceback +from pathlib import Path + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + +_GRAPH = Path(__file__).parent / "test_data" / "pick_and_place_maple_table_env_graph.yaml" + +# The spec pins the two On-table objects inside this box and the mug to a fixed point. +_POSITION_LIMITS_X = (0.55, 0.7) +_POSITION_LIMITS_Y = (-0.4, -0.1) +_MUG_TARGET_XYZ = (0.65, 0.25, 0.85) +_TOLERANCE_M = 0.05 + + +def _local_position(env, scene_name: str) -> list[float]: + """Return an object's (x, y, z) position in the env-local frame.""" + import warp as wp + + asset = env.unwrapped.scene[scene_name] + pos_world = wp.to_torch(asset.data.root_pos_w)[0] + env_origin = env.unwrapped.scene.env_origins[0].to(pos_world.device) + return (pos_world - env_origin).tolist() + + +def _test_yaml_spec_placement_satisfies_relations(simulation_app) -> bool: + """Solved placement must satisfy the spec's position_limits and at_position relations.""" + from isaaclab_arena.cli.isaaclab_arena_cli import arena_env_builder_cfg_from_argparse, get_isaaclab_arena_cli_parser + from isaaclab_arena.environment_spec.arena_env_graph_spec import ArenaEnvGraphSpec + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + + env = None + try: + spec = ArenaEnvGraphSpec.from_yaml(_GRAPH) + arena_env = spec.to_arena_env() + args_cli = get_isaaclab_arena_cli_parser().parse_args(["--num_envs", "1"]) + env = ArenaEnvBuilder(arena_env, arena_env_builder_cfg_from_argparse(args_cli)).make_registered() + env.reset() + + # Both On-table objects share one position_limits box; solving must land them inside it. + for name in ("rubiks_cube_hot3d_robolab", "bowl_ycb_robolab"): + x, y, _ = _local_position(env, name) + assert ( + _POSITION_LIMITS_X[0] - _TOLERANCE_M <= x <= _POSITION_LIMITS_X[1] + _TOLERANCE_M + ), f"{name} x={x:.3f} escaped position_limits {_POSITION_LIMITS_X}" + assert ( + _POSITION_LIMITS_Y[0] - _TOLERANCE_M <= y <= _POSITION_LIMITS_Y[1] + _TOLERANCE_M + ), f"{name} y={y:.3f} escaped position_limits {_POSITION_LIMITS_Y}" + + # The mug is pinned by at_position, so it lands on the requested point. + mug = _local_position(env, "mug_ycb_robolab") + for got, want, axis in zip(mug, _MUG_TARGET_XYZ, "xyz"): + assert abs(got - want) < _TOLERANCE_M, f"mug {axis}={got:.3f} missed at_position target {want}" + + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + finally: + if env is not None: + env.close() + + return True + + +def test_yaml_spec_placement_satisfies_relations(): + """Pytest entry point: YAML spec -> built env -> solved placement honors relations.""" + result = run_simulation_app_function(_test_yaml_spec_placement_satisfies_relations, headless=True) + assert result, f"Test {test_yaml_spec_placement_satisfies_relations.__name__} failed" + + +if __name__ == "__main__": + test_yaml_spec_placement_satisfies_relations() diff --git a/isaaclab_arena/tests/test_relation_solver_background_collision.py b/isaaclab_arena/tests/test_relation_solver_background_collision.py index 052864755c..b2abe1660e 100644 --- a/isaaclab_arena/tests/test_relation_solver_background_collision.py +++ b/isaaclab_arena/tests/test_relation_solver_background_collision.py @@ -16,8 +16,8 @@ def _make_desk(): """Anchor desk, 2m x 1m wide so a box has room to relocate along X away from the obstacle.""" - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.relations import IsAnchor + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -32,7 +32,7 @@ def _make_desk(): def _make_box(name: str = "box"): """A 0.3m cube to place On the desk (smaller than each desk half so a valid spot exists).""" - from isaaclab_arena.assets.dummy_object import DummyObject + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox return DummyObject( @@ -42,12 +42,8 @@ def _make_box(name: str = "box"): def _make_background(): - """Tall obstacle on the desk's left strip (x in [0, 0.5]). - - Narrower than the desk so a straddling box has a non-zero escape gradient: the - overlap-volume loss is flat when one box is fully enclosed by the other. - """ - from isaaclab_arena.assets.dummy_object import DummyObject + """Return an obstacle narrower than the desk to preserve a non-zero escape gradient.""" + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -63,7 +59,7 @@ def _mesh_box(name: str, extents: tuple[float, float, float], position: tuple[fl """Dummy object with a box collision mesh and fixed pose.""" import trimesh - from isaaclab_arena.assets.dummy_object import DummyObject + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -537,7 +533,7 @@ def test_arena_env_builder_forwards_background_collisions_by_default(monkeypatch from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.relation_solver_params import CollisionMode, RelationSolverParams - objects_with_relations = [object()] + objects_with_relations = [SimpleNamespace(name="object")] background_collision = object() calls = {} @@ -548,7 +544,11 @@ def get_objects_with_relations(self): return objects_with_relations def fake_solve_and_apply_relation_placement( - objects, num_envs, placer_params, collision_objects=None, scene_assets=None + objects, + num_envs, + placer_params, + collision_objects=None, + scene_assets=None, ): calls["objects"] = objects calls["num_envs"] = num_envs @@ -559,7 +559,7 @@ def fake_solve_and_apply_relation_placement( monkeypatch.setattr(builder_module, "solve_and_apply_relation_placement", fake_solve_and_apply_relation_placement) placer_params = ObjectPlacerParams(solver_params=RelationSolverParams(collision_mode=CollisionMode.MESH)) - arena_env = SimpleNamespace(scene=Scene(), placer_params=placer_params, embodiment=None) + arena_env = SimpleNamespace(scene=Scene(), embodiment=None, placer_params=placer_params) builder = ArenaEnvBuilder(arena_env, ArenaEnvBuilderCfg(num_envs=2)) builder._solve_relations() @@ -589,14 +589,18 @@ def get_objects_with_relations(self): return [] def fake_solve_and_apply_relation_placement( - objects, num_envs, placer_params, collision_objects=None, scene_assets=None + objects, + num_envs, + placer_params, + collision_objects=None, + scene_assets=None, ): calls["objects"] = objects calls["scene_assets"] = list(scene_assets) calls["collision_objects"] = collision_objects monkeypatch.setattr(builder_module, "solve_and_apply_relation_placement", fake_solve_and_apply_relation_placement) - arena_env = SimpleNamespace(scene=Scene(), placer_params=None, embodiment=None) + arena_env = SimpleNamespace(scene=Scene(), embodiment=None, placer_params=None) builder = ArenaEnvBuilder(arena_env, ArenaEnvBuilderCfg()) builder._solve_relations() @@ -606,14 +610,48 @@ def fake_solve_and_apply_relation_placement( assert calls["collision_objects"] is None +def test_arena_env_builder_includes_embodiment_relations(monkeypatch): + from types import SimpleNamespace + + import isaaclab_arena.environments.arena_env_builder as builder_module + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.arena_env_builder_cfg import ArenaEnvBuilderCfg + + calls = {} + + class Scene: + assets = {} + + def get_objects_with_relations(self): + return [] + + class Embodiment: + name = "droid" + + def get_relations(self): + return [object()] + + def fake_solve_and_apply_relation_placement(*args, **kwargs): + calls.update(kwargs) + calls["objects"] = args[0] + + monkeypatch.setattr(builder_module, "solve_and_apply_relation_placement", fake_solve_and_apply_relation_placement) + embodiment = Embodiment() + arena_env = SimpleNamespace(scene=Scene(), embodiment=embodiment, placer_params=None) + + ArenaEnvBuilder(arena_env, ArenaEnvBuilderCfg())._solve_relations() + + assert calls["objects"] == [embodiment] + + def test_relation_placement_includes_background_mesh_for_object_mesh_override(monkeypatch): """Object-level MESH override enables aggregate background meshes.""" import isaaclab_arena.environments.relation_solver_interface as interface_module - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.relation_solver_params import CollisionMode, RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox mesh_object = DummyObject( @@ -652,11 +690,11 @@ def test_relation_placement_includes_background_mesh_for_background_override(mon """A passive Background can opt into mesh collision when the solver default is BBOX.""" import isaaclab_arena.environments.relation_solver_interface as interface_module from isaaclab_arena.assets.background import Background - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.relation_solver_params import CollisionMode, RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox background = Background.__new__(Background) @@ -698,11 +736,11 @@ def test_relation_placement_skips_background_mesh_for_default_bbox(monkeypatch): """Default BBOX mode uses individual passive objects, not aggregate whole-scene meshes.""" import isaaclab_arena.environments.relation_solver_interface as interface_module from isaaclab_arena.assets.background import Background - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.relation_solver_params import CollisionMode, RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox background = Background.__new__(Background) diff --git a/isaaclab_arena/tests/test_relation_solver_embodiment.py b/isaaclab_arena/tests/test_relation_solver_embodiment.py new file mode 100644 index 0000000000..9fc2484097 --- /dev/null +++ b/isaaclab_arena/tests/test_relation_solver_embodiment.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026, 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 + +"""Relation placement tests for embodiments.""" + +import torch + +from isaaclab_arena.relations.object_placer import ObjectPlacer +from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams +from isaaclab_arena.relations.relations import IsAnchor, On +from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment +from isaaclab_arena.tests.dummy_object import DummyObject +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.pose import Pose, PosePerEnv + + +def _make_floor_and_robot(): + floor = DummyObject( + name="floor", + bounding_box=AxisAlignedBoundingBox( + min_point=(-2.0, -2.0, -0.05), + max_point=(2.0, 2.0, 0.0), + ), + initial_pose=Pose.identity(), + ) + floor.add_relation(IsAnchor()) + robot = DummyEmbodiment( + name="robot", + bounding_box=AxisAlignedBoundingBox( + min_point=(-0.2, -0.2, 0.0), + max_point=(0.2, 0.2, 1.2), + ), + ) + robot.add_relation(On(floor, clearance_m=0.0)) + return floor, robot + + +def test_relation_solver_places_embodiment(): + floor, robot = _make_floor_and_robot() + + result = ObjectPlacer(ObjectPlacerParams(placement_seed=3)).place([floor, robot])[0] + + assert result.success + assert robot in result.positions + assert robot.get_initial_pose() is not None + + +def test_batched_embodiment_placement_stores_per_env_poses(): + floor, robot = _make_floor_and_robot() + + results = ObjectPlacer(ObjectPlacerParams(placement_seed=3)).place([floor, robot], num_envs=2) + + assert len(results) == 2 + initial_pose = robot.get_initial_pose() + assert isinstance(initial_pose, PosePerEnv) + assert len(initial_pose.poses) == 2 + + +def test_world_bounding_box_applies_positive_quarter_turn(): + asset = DummyObject( + name="box", + bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 1.0, 1.0)), + initial_pose=Pose( + position_xyz=(3.0, 4.0, 0.0), + rotation_xyzw=(0.0, 0.0, 2**-0.5, 2**-0.5), + ), + ) + + world_bbox = asset.get_world_bounding_box() + + assert torch.allclose(world_bbox.min_point, torch.tensor([[2.0, 4.0, 0.0]])) + assert torch.allclose(world_bbox.max_point, torch.tensor([[3.0, 6.0, 1.0]])) diff --git a/isaaclab_arena/tests/test_relation_solver_interface.py b/isaaclab_arena/tests/test_relation_solver_interface.py index 6a1728a8b0..7949352207 100644 --- a/isaaclab_arena/tests/test_relation_solver_interface.py +++ b/isaaclab_arena/tests/test_relation_solver_interface.py @@ -5,10 +5,12 @@ """Tests for the relation placement orchestration API.""" +import pytest + def _make_desk(): - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.relations import IsAnchor + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -22,7 +24,7 @@ def _make_desk(): def _make_box(name: str = "box"): - from isaaclab_arena.assets.dummy_object import DummyObject + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox return DummyObject( @@ -60,6 +62,28 @@ def test_solve_and_apply_relation_placement_with_no_objects_returns_empty_result assert placement_event_cfg is None +def test_solve_and_apply_relation_placement_requires_unique_asset_names(): + from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement + + with pytest.raises(AssertionError, match="names must be unique"): + solve_and_apply_relation_placement([_make_box(), _make_box()], num_envs=1) + + +def test_solve_and_apply_relation_placement_rejects_scene_name_collision(): + from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement + from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + + embodiment = DummyEmbodiment( + name="droid", + scene_name="robot", + bounding_box=AxisAlignedBoundingBox(min_point=(-0.2, -0.2, 0.0), max_point=(0.2, 0.2, 1.0)), + ) + + with pytest.raises(AssertionError, match="duplicate scene keys"): + solve_and_apply_relation_placement([_make_box("robot"), embodiment], num_envs=1) + + def test_solve_and_apply_relation_placement_with_only_anchors_returns_no_reset_event(): from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams @@ -98,43 +122,76 @@ def test_static_solve_and_apply_relation_placement_reuses_object_only_placement( assert len(initial_pose.poses) == 2 -def test_dynamic_spawn_pose_skips_objects_missing_from_fallback_layout(): +def test_dynamic_spawn_pose_rejects_layout_missing_non_anchor(): from isaaclab_arena.environments.relation_solver_interface import _apply_dynamic_spawn_pose desk = _make_desk() box = _make_box() placement_pool = _FakePlacementPool([_fallback_layout(positions={})]) - _apply_dynamic_spawn_pose( - objects=[desk, box], - placement_pool=placement_pool, - anchor_objects_set={desk}, - ) - - assert box.get_initial_pose() is None + with pytest.raises(AssertionError, match="missing non-anchor asset 'box'"): + _apply_dynamic_spawn_pose( + assets=[desk, box], + placement_pool=placement_pool, + anchor_assets={desk}, + ) -def test_dynamic_spawn_pose_event_params_use_runtime_objects(): +def test_dynamic_spawn_pose_event_params_use_runtime_assets(): from isaaclab_arena.environments.relation_solver_interface import _apply_dynamic_spawn_pose desk = _make_desk() box = _make_box() - placement_pool = _FakePlacementPool([_fallback_layout(positions={})]) + placement_pool = _FakePlacementPool([_fallback_layout(positions={box: (0.1, 0.2, 0.3)})]) event_cfg = _apply_dynamic_spawn_pose( - objects=[desk, box], + assets=[desk, box], placement_pool=placement_pool, - anchor_objects_set={desk}, + anchor_assets={desk}, ) - assert [obj.name for obj in event_cfg.params["objects"]] == ["desk", "box"] + assert [asset.name for asset in event_cfg.params["assets"]] == ["desk", "box"] assert "placement_pool" in event_cfg.params -def test_static_initial_poses_skip_object_when_any_layout_is_missing_position(capsys): - from isaaclab_arena.environments.relation_solver_interface import _apply_static_initial_poses +def test_static_embodiment_placement_stores_per_env_poses(): + from isaaclab_arena.environments.relation_solver_interface import _apply_relation_placement_result + from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams + from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import PosePerEnv + desk = _make_desk() + robot = DummyEmbodiment( + name="robot", + bounding_box=AxisAlignedBoundingBox( + min_point=(-0.2, -0.2, 0.0), + max_point=(0.2, 0.2, 1.0), + ), + ) + layouts = [ + _fallback_layout(positions={robot: (0.1, 0.2, 0.0)}), + _fallback_layout(positions={robot: (0.3, 0.4, 0.0)}), + ] + + event_cfg = _apply_relation_placement_result( + assets=[desk, robot], + placer_params=ObjectPlacerParams(resolve_on_reset=False), + placement_pool=_FakePlacementPool(layouts), + num_envs=2, + ) + + # Embodiments now store their solved pose per env like objects, so no coordinated reset event. + assert event_cfg is None + initial_pose = robot.get_initial_pose() + assert isinstance(initial_pose, PosePerEnv) + assert initial_pose.poses[0].position_xyz == (0.1, 0.2, 0.0) + assert initial_pose.poses[1].position_xyz == (0.3, 0.4, 0.0) + + +def test_static_initial_poses_reject_layout_missing_non_anchor(): + from isaaclab_arena.environments.relation_solver_interface import _apply_static_initial_poses + desk = _make_desk() missing_box = _make_box("missing_box") placed_box = _make_box("placed_box") @@ -143,16 +200,10 @@ def test_static_initial_poses_skip_object_when_any_layout_is_missing_position(ca _fallback_layout(positions={placed_box: (0.2, 0.0, 0.2)}), ]) - _apply_static_initial_poses( - objects=[desk, missing_box, placed_box], - placement_pool=placement_pool, - anchor_objects_set={desk}, - num_envs=2, - ) - captured = capsys.readouterr() - assert "missing_box" in captured.out - - assert missing_box.get_initial_pose() is None - placed_box_initial_pose = placed_box.get_initial_pose() - assert isinstance(placed_box_initial_pose, PosePerEnv) - assert len(placed_box_initial_pose.poses) == 2 + with pytest.raises(AssertionError, match="missing non-anchor asset 'missing_box'"): + _apply_static_initial_poses( + assets=[desk, missing_box, placed_box], + placement_pool=placement_pool, + anchor_assets={desk}, + num_envs=2, + ) diff --git a/isaaclab_arena/tests/test_validate_placement.py b/isaaclab_arena/tests/test_validate_placement.py index 160a6be607..0f8f59d340 100644 --- a/isaaclab_arena/tests/test_validate_placement.py +++ b/isaaclab_arena/tests/test_validate_placement.py @@ -8,12 +8,12 @@ import math import torch -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.object_placer import ObjectPlacer from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.placement_validation import PlacementCheck from isaaclab_arena.relations.placement_validators import NextToValidator, NotNextToValidator, OnRelationValidator from isaaclab_arena.relations.relations import NextTo, NotNextTo, On, RotateAroundSolution, Side +from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox diff --git a/isaaclab_arena/utils/scene_pose_writes.py b/isaaclab_arena/utils/scene_pose_writes.py new file mode 100644 index 0000000000..a5a9b0f2cf --- /dev/null +++ b/isaaclab_arena/utils/scene_pose_writes.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026, 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 + +"""Shared helpers for writing scene root poses into the sim.""" + +from __future__ import annotations + +import torch +from typing import Any + + +def write_scene_root_poses_to_sim( + scene_asset: Any, + scene_name: str, + pose_tensor: torch.Tensor, + env_ids: torch.Tensor, + device: torch.device, +) -> None: + """Write world-frame root poses for articulations, rigid bodies, or AssetBase extras. + + Args: + scene_asset: Scene entity to write (articulation, rigid object, or XForm prim view). + scene_name: Isaac Lab scene key, used only for error messages. + pose_tensor: ``(N, 7)`` poses in world frame with env origins already applied. + env_ids: Environment indices being written. + device: Torch device for zero-velocity tensors on the articulation path. + """ + num_envs = pose_tensor.shape[0] + zero_velocity = torch.zeros(num_envs, 6, device=device) + + write_root_pose = getattr(scene_asset, "write_root_pose_to_sim", None) + if write_root_pose is not None: + write_root_pose(pose_tensor, env_ids=env_ids) + scene_asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_ids) + return + + set_world_poses = getattr(scene_asset, "set_world_poses", None) + if set_world_poses is not None: + set_world_poses( + positions=pose_tensor[:, :3], + orientations=pose_tensor[:, 3:7], + indices=env_ids.detach().cpu(), + ) + return + + assert False, f"Scene asset '{scene_name}' does not support root pose writes" diff --git a/isaaclab_arena_curobo/ik_reachability_validator.py b/isaaclab_arena_curobo/ik_reachability_validator.py index b001001e65..3ac8db28a7 100644 --- a/isaaclab_arena_curobo/ik_reachability_validator.py +++ b/isaaclab_arena_curobo/ik_reachability_validator.py @@ -14,7 +14,7 @@ import torch from typing import TYPE_CHECKING -from isaaclab_arena.relations.placement_events import get_base_rotation_per_object +from isaaclab_arena.relations.placement_events import get_base_rotation_per_asset from isaaclab_arena.relations.placement_validation import PlacementCheck from isaaclab_arena.relations.placement_validator_registry import register_validator from isaaclab_arena.relations.placement_validators import PlacementValidator @@ -112,7 +112,7 @@ def _validate( """ objects = list(positions.keys()) anchors = set(get_anchor_objects(objects)) - base_rotations = get_base_rotation_per_object(objects) + base_rotations = get_base_rotation_per_asset(objects) world_poses = { obj: get_object_world_pose_from_layout(positions, orientations, obj, base_rotations) for obj in objects diff --git a/isaaclab_arena_curobo/placement_pool_ik_validation.py b/isaaclab_arena_curobo/placement_pool_ik_validation.py index 522b2ed81a..98a99102cc 100644 --- a/isaaclab_arena_curobo/placement_pool_ik_validation.py +++ b/isaaclab_arena_curobo/placement_pool_ik_validation.py @@ -16,8 +16,8 @@ from typing import TYPE_CHECKING from isaaclab_arena.relations.placement_events import ( - get_base_rotation_per_object, - get_movable_object_names, + get_base_rotation_per_asset, + get_movable_asset_names, get_placement_pool, write_layout_to_sim, ) @@ -98,12 +98,12 @@ def validate_pool_ik( objects = placement_pool.objects anchor_objects_set = set(get_anchor_objects(objects)) - base_rotations = get_base_rotation_per_object(objects) + base_rotations = get_base_rotation_per_asset(objects) # TODO(xinjieyao, 2026-06-29): Expose the objects-to-reach as an interface # rather than checking every movable object. Default it to the task's pickup object(s) # then validate only those grasps. For now we batch-check all # movable objects. - movable_object_names = get_movable_object_names(objects, anchor_objects_set) + movable_object_names = get_movable_asset_names(objects, anchor_objects_set) layouts_per_env = placement_pool.layouts_per_env() num_envs = min(len(layouts_per_env), env.unwrapped.num_envs) diff --git a/isaaclab_arena_curobo/tests/test_ik_reachability_validator.py b/isaaclab_arena_curobo/tests/test_ik_reachability_validator.py index bec1c986df..e751d084a1 100644 --- a/isaaclab_arena_curobo/tests/test_ik_reachability_validator.py +++ b/isaaclab_arena_curobo/tests/test_ik_reachability_validator.py @@ -22,11 +22,11 @@ def _make_desk_box_pool(num_envs: int = 1, min_layouts_per_env: int = 2): """Build a small valid desk (anchor) + box (On desk) pool and return it.""" - from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relations import IsAnchor, On + from isaaclab_arena.tests.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose diff --git a/isaaclab_arena_curobo/utils/planner_utils.py b/isaaclab_arena_curobo/utils/planner_utils.py index 6f1f35214a..7650c9e1f0 100644 --- a/isaaclab_arena_curobo/utils/planner_utils.py +++ b/isaaclab_arena_curobo/utils/planner_utils.py @@ -141,7 +141,7 @@ def make_curobo_planner( debug_planner: Enable cuRobo planner debug output. """ if robot_scene_name is None: - robot_scene_name = embodiment.get_embodiment_name_in_scene() + robot_scene_name = embodiment.get_scene_name() planner_cfg = make_planner_cfg(embodiment, debug_planner=debug_planner) # cuRobo-Lab's MotionGen/collision world is single-env only for now. planner = CuroboPlanner(env=env, robot=env.scene[robot_scene_name], config=planner_cfg, env_id=env_id) diff --git a/isaaclab_arena_environments/droid_pick_and_place_lightwheel_kitchen.yaml b/isaaclab_arena_environments/droid_pick_and_place_lightwheel_kitchen.yaml new file mode 100644 index 0000000000..08740bb94d --- /dev/null +++ b/isaaclab_arena_environments/droid_pick_and_place_lightwheel_kitchen.yaml @@ -0,0 +1,71 @@ +# Copyright (c) 2026, 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 + +env_name: droid_pick_mustard_to_bowl +embodiment: + id: droid + registry_name: droid_abs_joint_pos + params: + stand_height_m: 0.8 +background: + id: kitchen + registry_name: lightwheel_robocasa_kitchen + params: {} +objects: +- id: mustard_bottle + registry_name: mustard_bottle_hope_robolab + params: {} +- id: bowl + registry_name: bowl_ycb_robolab + params: {} +object_references: +- id: right_counter_top + parent_id: kitchen + prim_path: counter_main_main_group/top_geometry_right + object_type: base + params: {} +- id: floor + parent_id: kitchen + prim_path: floor_room/geometry + object_type: base + params: {} +relations: +- kind: is_anchor + subject: floor + params: {} +- kind: is_anchor + subject: right_counter_top + params: {} +- kind: 'on' + subject: droid + reference: floor + params: {} +- kind: next_to + subject: droid + reference: right_counter_top + params: + side: negative_y + distance_m: 0.01 +- kind: rotate_around_solution + subject: droid + params: + yaw_rad: 1.57 +- kind: 'on' + subject: mustard_bottle + reference: right_counter_top + params: {} +- kind: 'on' + subject: bowl + reference: right_counter_top + params: {} +task: + composition: atomic + description: pick up the mustard bottle and place it in the bowl + subtasks: + - kind: PickAndPlaceTask + params: + pick_up_object: mustard_bottle + destination_location: bowl + background_scene: kitchen diff --git a/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py b/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py index 66826f5a05..398f916d4c 100644 --- a/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py +++ b/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py @@ -4,13 +4,12 @@ # SPDX-License-Identifier: Apache-2.0 # pyright: reportArgumentType=false -# ^^^ Suppress type errors for DummyObject → Object (duck typing works at runtime) +# ^^^ Suppress type errors for ExampleObject → Object (duck typing works at runtime) """Example notebook demonstrating the ObjectPlacer class without IsaacSim dependencies.""" # %% -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.object_placer import ObjectPlacer from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.relation_solver import RelationSolver @@ -18,6 +17,7 @@ from isaaclab_arena.relations.relations import IsAnchor, NextTo, On, Side, get_anchor_objects from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose +from isaaclab_arena_examples.relations.example_object import ExampleObject from isaaclab_arena_examples.relations.relation_solver_visualizer import RelationSolverVisualizer @@ -25,31 +25,31 @@ def run_dummy_object_placer_demo(): """Run the ObjectPlacer demo with dummy objects and a single anchor.""" # Create objects with bounding boxes - desk = DummyObject( + desk = ExampleObject( name="desk", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)) ) # Central object on the desk - center_box = DummyObject( + center_box = ExampleObject( name="center_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.15)) ) # Objects placed on each side of center_box - right_box = DummyObject( + right_box = ExampleObject( name="right_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.1)) ) - left_box = DummyObject( + left_box = ExampleObject( name="left_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.1)) ) - front_box = DummyObject( + front_box = ExampleObject( name="front_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.08)) ) - back_box = DummyObject( + back_box = ExampleObject( name="back_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.08)) ) # Box on top of center_box - top_box = DummyObject( + top_box = ExampleObject( name="top_box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.08)) ) @@ -105,23 +105,23 @@ def run_dummy_object_placer_demo(): def run_dummy_multi_anchor_demo(): """Demonstrate multiple anchors: objects placed relative to different fixed references.""" # Create anchor objects (fixed positions) - table = DummyObject( + table = ExampleObject( name="table", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 0.6, 0.75)), ) - chair = DummyObject( + chair = ExampleObject( name="chair", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.5, 0.5, 0.45)), ) - mug = DummyObject( + mug = ExampleObject( name="mug", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.1)), ) - book = DummyObject( + book = ExampleObject( name="book", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.15, 0.03)), ) - bin_obj = DummyObject( + bin_obj = ExampleObject( name="bin", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.4)), ) @@ -167,22 +167,22 @@ def run_dummy_multi_anchor_demo(): def run_dummy_no_collision_demo(): """Run RelationSolver with three boxes starting overlapping; animation shows them separate.""" # Create table (anchor) and three boxes - table = DummyObject( + table = ExampleObject( name="table", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.8, 0.6, 0.4)), ) table.add_relation(IsAnchor()) table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) - box_a = DummyObject( + box_a = ExampleObject( name="box_a", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.1)), ) - box_b = DummyObject( + box_b = ExampleObject( name="box_b", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.08)), ) - box_c = DummyObject( + box_c = ExampleObject( name="box_c", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.18, 0.1, 0.06)), ) diff --git a/isaaclab_arena_examples/relations/example_object.py b/isaaclab_arena_examples/relations/example_object.py new file mode 100644 index 0000000000..04478bf650 --- /dev/null +++ b/isaaclab_arena_examples/relations/example_object.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026, 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 __future__ import annotations + +from isaaclab_arena.relations.placement_asset import PlacementAsset +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + + +class ExampleObject(PlacementAsset): + """Box-shaped placement asset for the relation-solver example notebooks, with no Isaac Sim dependency.""" + + def __init__(self, name: str, bounding_box: AxisAlignedBoundingBox): + super().__init__(name=name) + self._bounding_box = bounding_box + + def get_bounding_box(self) -> AxisAlignedBoundingBox: + return self._bounding_box diff --git a/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py b/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py index 6d52fbecd4..3c25634f59 100644 --- a/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py +++ b/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py @@ -17,20 +17,20 @@ import numpy as np from matplotlib.patches import Rectangle -from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.relations.relation_solver import RelationSolver from isaaclab_arena.relations.relation_solver_params import RelationSolverParams from isaaclab_arena.relations.relation_solver_state import RelationSolverState from isaaclab_arena.relations.relations import IsAnchor, NextTo, Side from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose +from isaaclab_arena_examples.relations.example_object import ExampleObject def create_loss_heatmap_2d( solver: RelationSolver, - anchor_object: DummyObject, - child: DummyObject, - all_objects: list[DummyObject], + anchor_object: ExampleObject, + child: ExampleObject, + all_objects: list[ExampleObject], grid_resolution=50, x_range=(-0.5, 2.0), y_range=(-0.5, 2.0), @@ -154,17 +154,17 @@ def run_visualization_demo(): distance_m = 0.1 # Create parent object - parent = DummyObject(name="parent", bounding_box=parent_bbox) + parent = ExampleObject(name="parent", bounding_box=parent_bbox) parent.set_initial_pose(Pose(position_xyz=parent_pos, rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) parent.add_relation(IsAnchor()) # Create first child - placed to the RIGHT of parent - child1 = DummyObject(name="child1", bounding_box=child_bbox) + child1 = ExampleObject(name="child1", bounding_box=child_bbox) child1.add_relation(NextTo(parent, side=Side.POSITIVE_X, distance_m=distance_m)) child1.set_initial_pose(Pose(position_xyz=(0.5, 0.0, 0.05), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) # Initial guess # Create second child - placed to the RIGHT of child1 (chained placement) - child2 = DummyObject(name="child2", bounding_box=child_bbox) + child2 = ExampleObject(name="child2", bounding_box=child_bbox) child2.add_relation(NextTo(child1, side=Side.POSITIVE_X, distance_m=distance_m)) child2.set_initial_pose(Pose(position_xyz=(0.8, 0.0, 0.05), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) # Initial guess