From 735b053f7cd62de95c3a825db44e3778237aadb5 Mon Sep 17 00:00:00 2001 From: zhx06 Date: Wed, 15 Jul 2026 10:21:11 -0700 Subject: [PATCH 01/19] introduce PlacementEntity Signed-off-by: zhx06 --- isaaclab_arena/assets/dummy_embodiment.py | 49 +++++++ isaaclab_arena/assets/dummy_object.py | 35 +---- isaaclab_arena/assets/object.py | 4 - isaaclab_arena/assets/object_base.py | 56 ++------ isaaclab_arena/embodiments/embodiment_base.py | 60 ++++++-- .../environments/arena_env_builder.py | 19 ++- .../environments/relation_solver_interface.py | 134 +++++++++++------- .../relations/bounding_box_helpers.py | 20 +-- isaaclab_arena/relations/mesh_pair_cache.py | 10 +- isaaclab_arena/relations/no_overlap_aabb.py | 12 +- isaaclab_arena/relations/no_overlap_mesh.py | 8 +- isaaclab_arena/relations/object_placer.py | 90 ++++++------ isaaclab_arena/relations/placement_entity.py | 83 +++++++++++ isaaclab_arena/relations/placement_events.py | 103 +++++++++++--- .../relations/placement_pool_validation.py | 35 ++++- isaaclab_arena/relations/placement_result.py | 12 +- .../relations/pooled_object_placer.py | 6 +- isaaclab_arena/relations/relation_solver.py | 22 +-- .../relations/relation_solver_state.py | 28 ++-- isaaclab_arena/relations/relations.py | 18 +-- isaaclab_arena/tests/test_placement_events.py | 55 +++++++ ...st_relation_solver_background_collision.py | 62 +++++++- .../tests/test_relation_solver_embodiment.py | 54 +++++++ .../tests/test_relation_solver_interface.py | 96 +++++++++---- 24 files changed, 757 insertions(+), 314 deletions(-) create mode 100644 isaaclab_arena/assets/dummy_embodiment.py create mode 100644 isaaclab_arena/relations/placement_entity.py create mode 100644 isaaclab_arena/tests/test_relation_solver_embodiment.py diff --git a/isaaclab_arena/assets/dummy_embodiment.py b/isaaclab_arena/assets/dummy_embodiment.py new file mode 100644 index 0000000000..651ba0bc66 --- /dev/null +++ b/isaaclab_arena/assets/dummy_embodiment.py @@ -0,0 +1,49 @@ +# 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 isaaclab_arena.relations.placement_entity import PlacementEntity +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, quaternion_to_90_deg_z_quarters +from isaaclab_arena.utils.pose import Pose + + +class DummyEmbodiment(PlacementEntity): + """Embodiment geometry without simulator dependencies.""" + + def __init__( + self, + name: str, + bounding_box: AxisAlignedBoundingBox, + initial_pose: Pose | None = None, + collision_mesh: trimesh.Trimesh | None = None, + ) -> None: + super().__init__(name=name, tags=["embodiment"]) + self.initial_pose = initial_pose + self.bounding_box = bounding_box + self.collision_mesh = collision_mesh + + def get_bounding_box(self) -> AxisAlignedBoundingBox: + """Return root-relative bounds.""" + return self.bounding_box + + def get_world_bounding_box(self) -> AxisAlignedBoundingBox: + """Return bounds transformed by the root pose.""" + 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_collision_mesh(self) -> trimesh.Trimesh | None: + """Return the configured collision mesh.""" + return self.collision_mesh + + def supports_per_env_initial_pose(self) -> bool: + """Return False because the dummy stores one root pose.""" + return False diff --git a/isaaclab_arena/assets/dummy_object.py b/isaaclab_arena/assets/dummy_object.py index 02f95daec7..ca65003b20 100644 --- a/isaaclab_arena/assets/dummy_object.py +++ b/isaaclab_arena/assets/dummy_object.py @@ -7,13 +7,13 @@ 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.relations.placement_entity import PlacementEntity +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.pose import Pose -class DummyObject: +class DummyObject(PlacementEntity): """Dummy object for testing purposes without Isaac Sim dependencies.""" def __init__( @@ -21,29 +21,16 @@ def __init__( name: str, bounding_box: AxisAlignedBoundingBox, initial_pose: Pose | None = None, - relations: list[RelationBase] = [], + relations: list[RelationBase] | None = None, collision_mesh: trimesh.Trimesh | None = None, **kwargs, ): - self.name = name + 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) + self.relations = list(relations or []) 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).""" @@ -62,19 +49,9 @@ def get_world_bounding_box(self) -> AxisAlignedBoundingBox: 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..fdeb669629 100644 --- a/isaaclab_arena/assets/object.py +++ b/isaaclab_arena/assets/object.py @@ -66,10 +66,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 diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py index 19abe3949b..e59e1fd06a 100644 --- a/isaaclab_arena/assets/object_base.py +++ b/isaaclab_arena/assets/object_base.py @@ -19,17 +19,13 @@ 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_entity import PlacementEntity 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 +36,7 @@ ] -class ObjectBase(Asset, ABC): +class ObjectBase(PlacementEntity, ABC): """Parent class for (spawnable) Object and ObjectReference.""" def __init__( @@ -57,33 +53,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.""" @@ -118,6 +90,17 @@ def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: self.object_cfg.init_state.rot = initial_pose.rotation_xyzw self.event_cfg = self._init_event_cfg() + def set_placement_initial_pose(self, pose: Pose) -> None: + """Set a solved root pose without configuring an independent reset.""" + self.initial_pose = pose + if self.object_cfg is not None: + self.object_cfg.init_state.pos = pose.position_xyz + self.object_cfg.init_state.rot = pose.rotation_xyzw + + def has_pose_reset_event(self) -> bool: + """Return whether another reset event controls the root pose.""" + 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 +164,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/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index e55aa47a96..74b03786e6 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -3,20 +3,26 @@ # # 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.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_entity import PlacementEntity +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, quaternion_to_90_deg_z_quarters 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(PlacementEntity): name: str | None = None tags: list[str] = ["embodiment"] @@ -50,14 +56,43 @@ def __init__( self.xr: Any | None = None self.termination_cfg: Any | None = None - def set_initial_pose(self, pose: Pose) -> None: + def get_bounding_box(self) -> AxisAlignedBoundingBox: + """Return root-relative bounds for the USD-authored articulation pose.""" + 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.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: Compute bounds at configured initial joint positions when joint_pos is not None. + return compute_local_bounding_box_from_usd(spawn.usd_path, scale) + + def get_world_bounding_box(self) -> AxisAlignedBoundingBox: + """Return bounds transformed by the configured root pose.""" + bounding_box = self.get_bounding_box() + if self.initial_pose is None: + return bounding_box + quarters = quaternion_to_90_deg_z_quarters(self.initial_pose.rotation_xyzw) + return bounding_box.rotated_90_around_z(quarters).translated(self.initial_pose.position_xyz) + + def get_collision_mesh(self) -> trimesh.Trimesh | None: + """Return no mesh because embodiment placement uses bounds.""" + + def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: + """Set the embodiment root pose.""" + assert isinstance(pose, Pose), "Embodiments require one root Pose" self.initial_pose = pose + def supports_per_env_initial_pose(self) -> bool: + """Return False because embodiment configs store one root pose.""" + return False + 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.robot 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: """Env-local robot base pose, resolved in order: the explicit ``initial_pose`` override if set, @@ -138,10 +173,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.robot 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: diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py index b5e2c306b3..bbaa7e62c6 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 entities 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``. @@ -86,7 +85,12 @@ def _solve_relations(self) -> None: * **False** — applies one layout per environment so per-object reset events restore the same layout every time. """ - objects_with_relations = self.arena_env.scene.get_objects_with_relations() + placement_entities = list(self.arena_env.scene.get_objects_with_relations()) + scene_entity_names = {entity.name: entity.name for entity in placement_entities} + embodiment = self.arena_env.embodiment + if embodiment is not None and embodiment.get_relations(): + placement_entities.append(embodiment) + scene_entity_names[embodiment.name] = embodiment.get_embodiment_name_in_scene() placer_params = self.arena_env.placer_params if placer_params is None: @@ -102,10 +106,11 @@ 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_entities, num_envs=self.cfg.num_envs, placer_params=placer_params, scene_assets=self.arena_env.scene.assets.values(), + scene_entity_names=scene_entity_names, ) def get_all_variations(self) -> dict[str, list[VariationBase]]: diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py index 8895f2f919..5e2b6e371e 100644 --- a/isaaclab_arena/environments/relation_solver_interface.py +++ b/isaaclab_arena/environments/relation_solver_interface.py @@ -6,24 +6,28 @@ from __future__ import annotations import copy -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING 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, + place_entities_from_layouts, + 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_entity import PlacementEntity + from isaaclab_arena.relations.placement_result import PlacementResult def _get_passive_collision_objects( @@ -37,16 +41,17 @@ def _get_passive_collision_objects( def solve_and_apply_relation_placement( - objects: list[ObjectBase], + objects: list[PlacementEntity], num_envs: int, placer_params: ObjectPlacerParams | None = None, collision_objects: list[CollisionObject] | None = None, scene_assets: Iterable[Asset | RigidObjectSet] | None = None, + scene_entity_names: Mapping[str, str] | None = None, ) -> EventTermCfg | None: """Solve relation placement and apply the result to object reset/static state. Args: - objects: Objects with spatial predicates that should be relation-solved. + objects: Entities 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. @@ -54,6 +59,7 @@ def solve_and_apply_relation_placement( or relation-constrained. scene_assets: Optional scene assets to scan for passive collision objects when collision_objects is not supplied. + scene_entity_names: Isaac Lab scene name for each placement entity. Returns: Reset event config to attach to the environment when placement should be @@ -63,6 +69,11 @@ def solve_and_apply_relation_placement( if not objects: print("No objects with relations found in scene. Skipping relation solving.") return None + entity_names = {obj.name for obj in objects} + assert len(entity_names) == len(objects), "Placement entity names must be unique" + if scene_entity_names is None: + scene_entity_names = {obj.name: obj.name for obj in objects} + assert set(scene_entity_names) == entity_names, "scene_entity_names must contain every placement entity name" if placer_params is None: placer_params = ObjectPlacerParams() @@ -81,8 +92,6 @@ def solve_and_apply_relation_placement( objects, 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, placer_params=placer_params, @@ -105,11 +114,12 @@ def solve_and_apply_relation_placement( placer_params=placer_params, placement_pool=placement_pool, num_envs=num_envs, + scene_entity_names=scene_entity_names, ) def _should_include_background_mesh( - objects: list[ObjectBase], + objects: list[PlacementEntity], scene_assets: Iterable[Asset | RigidObjectSet], default_collision_mode: CollisionMode, ) -> bool: @@ -127,10 +137,11 @@ def _should_include_background_mesh( def _apply_relation_placement_result( - objects: list[ObjectBase], + objects: list[PlacementEntity], placer_params: ObjectPlacerParams, placement_pool: PooledObjectPlacer, num_envs: int, + scene_entity_names: Mapping[str, str], ) -> EventTermCfg | None: """Apply selected layouts to object spawn state and build reset event config.""" anchor_objects_set = set(get_anchor_objects(objects)) @@ -146,41 +157,38 @@ def _apply_relation_placement_result( objects=objects, placement_pool=placement_pool, anchor_objects_set=anchor_objects_set, + scene_entity_names=scene_entity_names, ) - _apply_static_initial_poses( + if all(obj.supports_per_env_initial_pose() for obj in objects): + _apply_static_initial_poses( + objects=objects, + placement_pool=placement_pool, + anchor_objects_set=anchor_objects_set, + num_envs=num_envs, + ) + return None + return _apply_static_spawn_pose( objects=objects, placement_pool=placement_pool, anchor_objects_set=anchor_objects_set, num_envs=num_envs, + scene_entity_names=scene_entity_names, ) - return None def _apply_dynamic_spawn_pose( - objects: list[ObjectBase], + objects: list[PlacementEntity], placement_pool: PooledObjectPlacer, - anchor_objects_set: set[ObjectBase], + anchor_objects_set: set[PlacementEntity], + scene_entity_names: Mapping[str, str], ) -> 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 + _set_placement_initial_poses(objects, anchor_objects_set, layout) return EventTermCfg( func=solve_and_place_objects, @@ -188,14 +196,51 @@ def _apply_dynamic_spawn_pose( params={ "objects": objects, "placement_pool": placement_pool, + "scene_entity_names": scene_entity_names, }, ) +def _apply_static_spawn_pose( + objects: list[PlacementEntity], + placement_pool: PooledObjectPlacer, + anchor_objects_set: set[PlacementEntity], + num_envs: int, + scene_entity_names: Mapping[str, str], +) -> EventTermCfg: + """Return a reset event that restores one fixed layout per environment.""" + from isaaclab.managers import EventTermCfg + + layouts = placement_pool.sample_with_replacement(num_envs) + _set_placement_initial_poses(objects, anchor_objects_set, layouts[0]) + return EventTermCfg( + func=place_entities_from_layouts, + mode="reset", + params={ + "objects": objects, + "layouts": layouts, + "scene_entity_names": scene_entity_names, + }, + ) + + +def _set_placement_initial_poses( + objects: list[PlacementEntity], + anchor_objects_set: set[PlacementEntity], + layout: PlacementResult, +) -> None: + """Seed non-anchor entities from one placement layout.""" + for obj in objects: + if obj in anchor_objects_set: + continue + pose = get_pose_from_layout(obj, layout) + obj.set_placement_initial_pose(pose) + + def _apply_static_initial_poses( - objects: list[ObjectBase], + objects: list[PlacementEntity], placement_pool: PooledObjectPlacer, - anchor_objects_set: set[ObjectBase], + anchor_objects_set: set[PlacementEntity], num_envs: int, ) -> None: """Apply fixed per-environment poses for ``resolve_on_reset=False``.""" @@ -203,34 +248,17 @@ def _apply_static_initial_poses( for obj in objects: if obj in anchor_objects_set: 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(obj, layouts[env_idx]) for env_idx in range(num_envs)] + obj.set_initial_pose(PosePerEnv(poses=poses)) def _validate_no_conflicting_pose_reset_events( - objects: list[ObjectBase], - anchor_objects_set: set[ObjectBase], + objects: list[PlacementEntity], + anchor_objects_set: set[PlacementEntity], ) -> 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), ( + assert not (obj not in anchor_objects_set and obj.has_pose_reset_event()), ( f"Non-anchor object '{obj.name}' has an explicit pose-reset event. " "Relational solving should not be combined with explicit setting of " "poses on non-anchor objects." diff --git a/isaaclab_arena/relations/bounding_box_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py index 1428d52e5f..e3f92670cb 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 entities. """ 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_entity import PlacementEntity -def has_heterogeneous_objects(objects: list[ObjectBase]) -> bool: +def has_heterogeneous_objects(objects: list[PlacementEntity]) -> 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[PlacementEntity], 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: PlacementEntity, 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[PlacementEntity, 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[PlacementEntity, 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[PlacementEntity, 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[PlacementEntity, 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[PlacementEntity], 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..bf7a06ac4f 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_entity import PlacementEntity 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: PlacementEntity """Subject (sphere source) object.""" - obstacle: ObjectBase | CollisionObject + obstacle: PlacementEntity | 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[PlacementEntity] """(P,) subject (sphere source) object reference per pair.""" - pair_obstacle_objs: list[ObjectBase | CollisionObject] + pair_obstacle_objs: list[PlacementEntity | 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..37877b1bad 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_entity import PlacementEntity 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[PlacementEntity | 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: PlacementEntity, 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: PlacementEntity, + b: PlacementEntity, 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: PlacementEntity, 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..f9728be559 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_entity import PlacementEntity 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[PlacementEntity, 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[PlacementEntity | 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: PlacementEntity, 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..6e0275d039 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_entity import PlacementEntity 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[PlacementEntity, 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[PlacementEntity, 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[PlacementEntity], 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[PlacementEntity], 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[PlacementEntity], + ) -> tuple[set[PlacementEntity], 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[PlacementEntity], + anchor_objects_set: set[PlacementEntity], 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[PlacementEntity, tuple[float, float, float]]] = [] + orientations_per_candidate: list[dict[PlacementEntity, 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[PlacementEntity], + anchor_objects: set[PlacementEntity], + env_bboxes: dict[PlacementEntity, AxisAlignedBoundingBox], generator: torch.Generator | None = None, - ) -> dict[ObjectBase, tuple[float, float, float]]: + ) -> dict[PlacementEntity, 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[PlacementEntity, 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: PlacementEntity, + env_bboxes: dict[PlacementEntity, 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[PlacementEntity], + anchor_objects: set[PlacementEntity], generator: torch.Generator | None = None, - ) -> dict[ObjectBase, float]: + ) -> dict[PlacementEntity, 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[PlacementEntity, 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[PlacementEntity, tuple[float, float, float]]], + orientations_per_candidate: list[dict[PlacementEntity, 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[PlacementEntity], + candidate_bboxes: dict[PlacementEntity, AxisAlignedBoundingBox], + orientations_per_candidate: list[dict[PlacementEntity, float]], + ) -> dict[PlacementEntity, 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[PlacementEntity, 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[PlacementEntity, AxisAlignedBoundingBox], candidate_idx: int, - ) -> dict[ObjectBase, AxisAlignedBoundingBox]: + ) -> dict[PlacementEntity, 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: PlacementEntity, + anchor_objects: set[PlacementEntity], anchor_bbox: AxisAlignedBoundingBox, - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + env_bboxes: dict[PlacementEntity, 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: PlacementEntity, + anchor_objects: set[PlacementEntity], anchor_bbox: AxisAlignedBoundingBox, - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox], + env_bboxes: dict[PlacementEntity, 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[PlacementEntity, tuple[float, float, float]]], + orientations: list[dict[PlacementEntity, float]], + bboxes: list[dict[PlacementEntity, AxisAlignedBoundingBox]], collision_objects: list[CollisionObject], ) -> list[PlacementValidationResults]: """Run every enabled validator over all candidates and collect per-candidate results. @@ -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[PlacementEntity, tuple[float, float, float]]], + anchor_objects: set[PlacementEntity], + orientations_per_env: list[dict[PlacementEntity, float]], ) -> None: """Apply solved positions and orientations to non-anchor objects. @@ -731,6 +731,10 @@ def _yaw_delta(env_idx: int) -> float: else: obj.set_initial_pose(Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw)) else: + assert obj.supports_per_env_initial_pose(), ( + f"Placement entity '{obj.name}' cannot store per-environment poses. " + "Set apply_positions_to_objects=False and apply the results at reset." + ) poses = [ Pose( position_xyz=positions_per_env[env_idx][obj], diff --git a/isaaclab_arena/relations/placement_entity.py b/isaaclab_arena/relations/placement_entity.py new file mode 100644 index 0000000000..830dbf39c3 --- /dev/null +++ b/isaaclab_arena/relations/placement_entity.py @@ -0,0 +1,83 @@ +# 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 entities 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.pose import Pose, PosePerEnv, PoseRange + +if TYPE_CHECKING: + import trimesh + + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + + +class PlacementEntity(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 + # Mesh collision uses a convex hull when source geometry is not watertight. + self.repair_collision_mesh_non_watertight = True + + def add_relation(self, relation: RelationBase) -> None: + """Attach a relation to the entity.""" + self.relations.append(relation) + + def get_relations(self) -> list[RelationBase]: + """Return all relations attached to the entity.""" + 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 entity 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.""" + self.initial_pose = pose + + def set_placement_initial_pose(self, pose: Pose) -> None: + """Set a solved root pose without changing reset ownership.""" + self.set_initial_pose(pose) + + def has_pose_reset_event(self) -> bool: + """Return whether the entity owns a root-pose reset event.""" + return False + + def supports_per_env_initial_pose(self) -> bool: + """Return whether one configured pose may be stored per environment.""" + return True + + @abstractmethod + def get_bounding_box(self) -> AxisAlignedBoundingBox: + """Return root-relative axis-aligned bounds.""" + + @abstractmethod + def get_world_bounding_box(self) -> AxisAlignedBoundingBox: + """Return axis-aligned bounds transformed by the root pose.""" + + @abstractmethod + def get_collision_mesh(self) -> trimesh.Trimesh | None: + """Return a root-relative collision mesh when available.""" diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py index 7a4e330ec1..e8bb189fe0 100644 --- a/isaaclab_arena/relations/placement_events.py +++ b/isaaclab_arena/relations/placement_events.py @@ -6,6 +6,7 @@ from __future__ import annotations import torch +from collections.abc import Mapping from typing import TYPE_CHECKING from isaaclab_arena.relations.relations import RotateAroundSolution, get_anchor_objects @@ -16,7 +17,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv - from isaaclab_arena.assets.object_base import ObjectBase + from isaaclab_arena.relations.placement_entity import PlacementEntity from isaaclab_arena.relations.placement_result import PlacementResult from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer @@ -42,31 +43,54 @@ 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]: +def get_placement_scene_entity_names(env) -> Mapping[str, str] | None: + """Return the scene-name map stored on the placement reset event.""" + try: + term_cfg = env.unwrapped.event_manager.get_term_cfg(PLACEMENT_RESET_EVENT_NAME) + except ValueError: + return None + return term_cfg.params.get("scene_entity_names") + + +def get_rotation_xyzw(obj: PlacementEntity) -> 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) 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]]: +def get_base_rotation_per_object( + objects: list[PlacementEntity], +) -> dict[PlacementEntity, tuple[float, float, float, float]]: """Return the base rotation for each object.""" return {obj: get_rotation_xyzw(obj) for obj in objects} +def get_pose_from_layout(obj: PlacementEntity, layout: PlacementResult) -> Pose: + """Return an entity pose from a solved layout.""" + assert obj in layout.positions, f"Placement layout is missing non-anchor entity '{obj.name}'" + base_rotation = get_rotation_xyzw(obj) + marker_yaw = yaw_from_quat_xyzw(base_rotation) + total_yaw = layout.orientations.get(obj, marker_yaw) + rotation = rotate_quat_by_yaw(base_rotation, total_yaw - marker_yaw) + return Pose(position_xyz=layout.positions[obj], rotation_xyzw=rotation) + + def get_movable_object_names( - objects: list[ObjectBase], - anchor_objects_set: set[ObjectBase], + objects: list[PlacementEntity], + anchor_objects_set: set[PlacementEntity], + scene_entity_names: Mapping[str, str], ) -> 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 entities.""" + return [scene_entity_names[obj.name] for obj in objects if obj not in anchor_objects_set] 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_objects_set: set[PlacementEntity], + base_rotations: dict[PlacementEntity, tuple[float, float, float, float]], + scene_entity_names: Mapping[str, str] | None = None, ) -> None: """Write one env's solved layout into the sim. @@ -79,28 +103,34 @@ def write_layout_to_sim( 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. + scene_entity_names: Isaac Lab scene name for each placement entity. Entity + names are used directly when omitted. """ + if scene_entity_names is None: + scene_entity_names = {obj.name: obj.name for obj in base_rotations} 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(): + missing_entities = [ + obj.name for obj in base_rotations if obj not in anchor_objects_set and obj not in result.positions + ] + assert not missing_entities, f"Placement layout is missing non-anchor entities: {missing_entities}" + for obj in result.positions: if obj in anchor_objects_set: 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 = env.scene[scene_entity_names[obj.name]] + pose = get_pose_from_layout(obj, result) + pose_tensor = pose.to_tensor(device=env.device).unsqueeze(0) + pose_tensor[0, :3] += env.scene.env_origins[env_id, :] + asset.write_root_pose_to_sim(pose_tensor, env_ids=env_id_tensor) asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_id_tensor) def solve_and_place_objects( env: ManagerBasedEnv, env_ids: torch.Tensor | None, - objects: list[ObjectBase], + objects: list[PlacementEntity], placement_pool: PooledObjectPlacer, + scene_entity_names: Mapping[str, str], ) -> None: """Coordinated reset event that draws layouts from the pool and writes poses. @@ -113,6 +143,7 @@ def solve_and_place_objects( env_ids: 1-D tensor of environment indices being reset. objects: Objects participating in relation solving. placement_pool: Runtime pool of solved placement layouts. + scene_entity_names: Isaac Lab scene name for each placement entity. """ if env_ids is None or len(env_ids) == 0: return @@ -133,4 +164,36 @@ def solve_and_place_objects( 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) + write_layout_to_sim(env, cur_env, result, anchor_objects_set, base_rotations, scene_entity_names) + + +def place_entities_from_layouts( + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + objects: list[PlacementEntity], + layouts: list[PlacementResult], + scene_entity_names: Mapping[str, str], +) -> None: + """Restore one fixed placement layout per environment. + + Args: + env: The Isaac Lab environment. + env_ids: Environment indices to restore. + objects: Entities participating in relation solving. + layouts: Fixed layout indexed by environment. + scene_entity_names: Isaac Lab scene name for each placement entity. + """ + if env_ids is None or len(env_ids) == 0: + return + assert len(layouts) == env.scene.env_origins.shape[0], "Static layouts must match the scene environment count" + anchor_objects_set = set(get_anchor_objects(objects)) + base_rotations = get_base_rotation_per_object(objects) + for env_id in env_ids.tolist(): + write_layout_to_sim( + env, + env_id, + layouts[env_id], + anchor_objects_set, + base_rotations, + scene_entity_names, + ) diff --git a/isaaclab_arena/relations/placement_pool_validation.py b/isaaclab_arena/relations/placement_pool_validation.py index ee8f73aeb3..1e437e2330 100644 --- a/isaaclab_arena/relations/placement_pool_validation.py +++ b/isaaclab_arena/relations/placement_pool_validation.py @@ -5,6 +5,7 @@ from __future__ import annotations +from collections.abc import Mapping from typing import TYPE_CHECKING from isaaclab_arena.relations.physics_settle_params import PhysicsSettleParams @@ -12,6 +13,7 @@ get_base_rotation_per_object, get_movable_object_names, get_placement_pool, + get_placement_scene_entity_names, write_layout_to_sim, ) from isaaclab_arena.relations.placement_validation import PlacementCheck @@ -21,7 +23,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv - from isaaclab_arena.relations.object_base import ObjectBase + from isaaclab_arena.relations.placement_entity import PlacementEntity 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 @@ -33,7 +35,8 @@ def _write_layout_to_envs_for_episode_index( num_envs: int, episode_index: int, anchor_objects_set: set, - base_rotations: dict[ObjectBase, tuple[float, float, float, float]], + base_rotations: dict[PlacementEntity, tuple[float, float, float, float]], + scene_entity_names: Mapping[str, str], ) -> list[tuple[int, PlacementResult]]: """Write one layout per env for this episode; return the ``(env_id, layout)`` layouts written. @@ -45,7 +48,14 @@ 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_objects_set, + base_rotations, + scene_entity_names, + ) layouts_written.append((env_id, layout)) return layouts_written @@ -79,6 +89,7 @@ def validate_pool_layouts( placement_pool: PooledObjectPlacer | None = None, settle_params: PhysicsSettleParams | None = None, render: bool = False, + scene_entity_names: Mapping[str, str] | None = None, ) -> list[tuple[int, int, PlacementValidationResults]] | None: """Physics-validate every layout in a placement pool, recording the result on its validation results. @@ -92,12 +103,13 @@ def validate_pool_layouts( settle_params: Settle-check tuning params. Defaults to ``PhysicsSettleParams()`` when omitted. render: When True, render each settle step so the sweep is visible in the GUI. Defaults to False. + scene_entity_names: Isaac Lab scene name for each placement entity. Returns: ``(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 + placement_pool_from_event = placement_pool is None if placement_pool is None: placement_pool = get_placement_pool(env) if placement_pool is None: @@ -108,7 +120,12 @@ def validate_pool_layouts( 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) + if scene_entity_names is None: + scene_entity_names = get_placement_scene_entity_names(env) + if scene_entity_names is None: + assert not placement_pool_from_event, "Placement reset event is missing scene_entity_names" + scene_entity_names = {obj.name: obj.name for obj in objects} + movable_object_names = get_movable_object_names(objects, anchor_objects_set, scene_entity_names) # 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 +144,13 @@ 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_objects_set, + base_rotations, + scene_entity_names, ) 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..be9a56bc48 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_entity import PlacementEntity from isaaclab_arena.relations.placement_validation import PlacementValidationResults @dataclass class PlacementResult: - """Solved object layout for one environment.""" + """Solved entity 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[PlacementEntity, tuple[float, float, float]] + """Final positions for each entity.""" final_loss: float """Loss value of the final placement.""" @@ -29,8 +29,8 @@ 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[PlacementEntity, float] = field(default_factory=dict) + """Placement-computed world Z-yaws. Omitted entities retain their marker orientation.""" @property def success(self) -> bool: diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py index 17fbc9cddf..522bded01b 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_entity import PlacementEntity @dataclass @@ -70,7 +70,7 @@ class PooledObjectPlacer: def __init__( self, - objects: list[ObjectBase], + objects: list[PlacementEntity], 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[PlacementEntity]: """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..0ffc4cc3b9 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_entity import PlacementEntity 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[PlacementEntity, 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[PlacementEntity], + initial_positions: list[dict[PlacementEntity, tuple[float, float, float]]], + env_bboxes: dict[PlacementEntity, AxisAlignedBoundingBox] | None = None, env_bboxes_include_yaw: bool = False, - orientations: list[dict[ObjectBase, float]] | None = None, + orientations: list[dict[PlacementEntity, float]] | None = None, collision_objects: list[CollisionObject] | None = None, - ) -> list[dict[ObjectBase, tuple[float, float, float]]]: + ) -> list[dict[PlacementEntity, 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 entities to solve. Must include at least one entity 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[PlacementEntity]) -> 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: PlacementEntity, 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: PlacementEntity, 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..134b031336 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_entity import PlacementEntity 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[PlacementEntity], + initial_positions: list[dict[PlacementEntity, tuple[float, float, float]]], device: torch.device | None = None, - env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None, + env_bboxes: dict[PlacementEntity, 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 entities 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[PlacementEntity] = 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[PlacementEntity, 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[PlacementEntity | 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[PlacementEntity]: """List of optimizable objects (excludes anchors).""" return self._optimizable_objects @property - def anchor_objects(self) -> set[ObjectBase]: + def anchor_objects(self) -> set[PlacementEntity]: """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: PlacementEntity) -> 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: PlacementEntity | 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: PlacementEntity) -> 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[PlacementEntity, 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..045e73a60b 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_entity import PlacementEntity 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: PlacementEntity, objects: set[PlacementEntity]) -> 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: PlacementEntity, 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: PlacementEntity): """ 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: PlacementEntity, objects: set[PlacementEntity]) -> 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: PlacementEntity, 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: PlacementEntity, 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: PlacementEntity, 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[PlacementEntity]) -> list[PlacementEntity]: """Get all anchor objects from a list of objects. Anchor objects are marked with IsAnchor() relation and serve as diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index 496b2e4fb7..1d39f09d4d 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -164,6 +164,7 @@ def _solve_and_place_with_pool(env, env_ids, objects, pool): env_ids, objects=objects, placement_pool=pool, + scene_entity_names={obj.name: obj.name for obj in objects}, ) @@ -225,12 +226,61 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: torch.tensor([0]), objects=[desk, box1], placement_pool=Pool(), + scene_entity_names={"desk": "desk", "box1": "box1"}, ) assert "desk" not in env._assets env._assets["box1"].write_root_pose_to_sim.assert_called_once() +def test_static_layout_event_writes_embodiment_to_mapped_scene_asset(): + from isaaclab_arena.assets.dummy_embodiment import DummyEmbodiment + from isaaclab_arena.relations.placement_events import place_entities_from_layouts + from isaaclab_arena.relations.placement_result import PlacementResult + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + + robot = DummyEmbodiment( + name="droid", + bounding_box=AxisAlignedBoundingBox(min_point=(-0.2, -0.2, 0.0), max_point=(0.2, 0.2, 1.0)), + ) + layouts = [ + PlacementResult( + validation_results=_checklist(True), + positions={robot: (0.1, 0.2, 0.0)}, + final_loss=0.0, + attempts=1, + ), + PlacementResult( + validation_results=_checklist(True), + positions={robot: (0.4, 0.5, 0.0)}, + final_loss=0.0, + attempts=1, + ), + ] + env = _make_mock_env(num_envs=2) + + place_entities_from_layouts( + env, + torch.tensor([1]), + objects=[robot], + layouts=layouts, + scene_entity_names={"droid": "robot"}, + ) + + assert "droid" not in env._assets + pose = env._assets["robot"].write_root_pose_to_sim.call_args.args[0] + assert torch.allclose(pose[0, :3], torch.tensor([0.4, 0.5, 0.0])) + + with pytest.raises(AssertionError, match="Static layouts must match"): + place_entities_from_layouts( + env, + torch.tensor([0]), + objects=[robot], + layouts=layouts[:1], + scene_entity_names={"droid": "robot"}, + ) + + 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 @@ -581,6 +631,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_placement_initial_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 @@ -606,6 +660,7 @@ def sample_with_replacement(self, count: int): objects=[anchor, box], placement_pool=pool, anchor_objects_set={anchor}, + scene_entity_names={"desk": "desk", "box": "box"}, ) assert pool.sample_count == 1 diff --git a/isaaclab_arena/tests/test_relation_solver_background_collision.py b/isaaclab_arena/tests/test_relation_solver_background_collision.py index 052864755c..f27a59b778 100644 --- a/isaaclab_arena/tests/test_relation_solver_background_collision.py +++ b/isaaclab_arena/tests/test_relation_solver_background_collision.py @@ -537,7 +537,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,18 +548,24 @@ 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, + scene_entity_names=None, ): calls["objects"] = objects calls["num_envs"] = num_envs calls["placer_params"] = placer_params calls["scene_assets"] = list(scene_assets) calls["collision_objects"] = collision_objects + calls["scene_entity_names"] = scene_entity_names return "placement_event" 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() @@ -569,6 +575,7 @@ def fake_solve_and_apply_relation_placement( assert calls["placer_params"] is placer_params assert calls["scene_assets"] == [background_collision] assert calls["collision_objects"] is None + assert calls["scene_entity_names"]["object"] == "object" assert builder._placement_event_cfg == "placement_event" @@ -589,14 +596,20 @@ 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, + scene_entity_names=None, ): calls["objects"] = objects calls["scene_assets"] = list(scene_assets) calls["collision_objects"] = collision_objects + calls["scene_entity_names"] = scene_entity_names 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() @@ -604,6 +617,45 @@ def fake_solve_and_apply_relation_placement( assert calls["objects"] == [] assert calls["scene_assets"] == [] assert calls["collision_objects"] is None + assert calls["scene_entity_names"] == {} + + +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 get_embodiment_name_in_scene(self): + return "robot" + + 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] + assert calls["scene_entity_names"] == {"droid": "robot"} def test_relation_placement_includes_background_mesh_for_object_mesh_override(monkeypatch): 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..0b53debeaf --- /dev/null +++ b/isaaclab_arena/tests/test_relation_solver_embodiment.py @@ -0,0 +1,54 @@ +# 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 pytest + +from isaaclab_arena.assets.dummy_embodiment import DummyEmbodiment +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.relations import IsAnchor, On +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.pose import Pose + + +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_requires_runtime_application(): + floor, robot = _make_floor_and_robot() + + with pytest.raises(AssertionError, match="cannot store per-environment poses"): + ObjectPlacer(ObjectPlacerParams(placement_seed=3)).place([floor, robot], num_envs=2) diff --git a/isaaclab_arena/tests/test_relation_solver_interface.py b/isaaclab_arena/tests/test_relation_solver_interface.py index 6a1728a8b0..1f5609f6b8 100644 --- a/isaaclab_arena/tests/test_relation_solver_interface.py +++ b/isaaclab_arena/tests/test_relation_solver_interface.py @@ -5,6 +5,8 @@ """Tests for the relation placement orchestration API.""" +import pytest + def _make_desk(): from isaaclab_arena.assets.dummy_object import DummyObject @@ -60,6 +62,20 @@ 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_entity_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_requires_complete_scene_name_map(): + from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement + + with pytest.raises(AssertionError, match="must contain every placement entity"): + solve_and_apply_relation_placement([_make_desk()], num_envs=1, scene_entity_names={}) + + 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,20 +114,20 @@ 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 entity 'box'"): + _apply_dynamic_spawn_pose( + objects=[desk, box], + placement_pool=placement_pool, + anchor_objects_set={desk}, + scene_entity_names={"desk": "desk", "box": "box"}, + ) def test_dynamic_spawn_pose_event_params_use_runtime_objects(): @@ -119,21 +135,59 @@ def test_dynamic_spawn_pose_event_params_use_runtime_objects(): 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], placement_pool=placement_pool, anchor_objects_set={desk}, + scene_entity_names={"desk": "desk", "box": "box"}, ) assert [obj.name for obj in event_cfg.params["objects"]] == ["desk", "box"] assert "placement_pool" in event_cfg.params + assert event_cfg.params["scene_entity_names"]["box"] == "box" + +def test_static_embodiment_placement_uses_coordinated_reset(): + from isaaclab_arena.assets.dummy_embodiment import DummyEmbodiment + from isaaclab_arena.environments.relation_solver_interface import _apply_relation_placement_result + from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.pose import Pose + + 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( + objects=[desk, robot], + placer_params=ObjectPlacerParams(resolve_on_reset=False), + placement_pool=_FakePlacementPool(layouts), + num_envs=2, + scene_entity_names={"desk": "desk", "robot": "robot"}, + ) -def test_static_initial_poses_skip_object_when_any_layout_is_missing_position(capsys): + initial_pose = robot.get_initial_pose() + assert isinstance(initial_pose, Pose) + assert initial_pose.position_xyz == (0.1, 0.2, 0.0) + assert event_cfg is not None + runtime_robot = event_cfg.params["objects"][1] + assert runtime_robot in event_cfg.params["layouts"][0].positions + assert event_cfg.params["layouts"][1].positions[runtime_robot] == (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 - from isaaclab_arena.utils.pose import PosePerEnv desk = _make_desk() missing_box = _make_box("missing_box") @@ -143,16 +197,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 entity 'missing_box'"): + _apply_static_initial_poses( + objects=[desk, missing_box, placed_box], + placement_pool=placement_pool, + anchor_objects_set={desk}, + num_envs=2, + ) From 535dded0630600675613052f1dac922647ed635c Mon Sep 17 00:00:00 2001 From: zhx06 Date: Tue, 21 Jul 2026 14:34:22 -0700 Subject: [PATCH 02/19] address review comments Signed-off-by: zhx06 --- isaaclab_arena/assets/asset.py | 4 + isaaclab_arena/assets/dummy_object.py | 16 +- isaaclab_arena/assets/object.py | 23 +-- isaaclab_arena/assets/object_base.py | 15 +- isaaclab_arena/embodiments/embodiment_base.py | 39 ++--- .../environments/arena_env_builder.py | 16 +- .../environments/relation_solver_interface.py | 140 +++++++++--------- .../relations/bounding_box_helpers.py | 20 +-- isaaclab_arena/relations/mesh_pair_cache.py | 10 +- isaaclab_arena/relations/no_overlap_aabb.py | 12 +- isaaclab_arena/relations/no_overlap_mesh.py | 8 +- isaaclab_arena/relations/object_placer.py | 86 +++++------ ...placement_entity.py => placement_asset.py} | 38 +++-- isaaclab_arena/relations/placement_events.py | 113 ++++++-------- .../relations/placement_pool_validation.py | 35 ++--- isaaclab_arena/relations/placement_result.py | 17 +-- .../relations/placement_validators.py | 98 ++++++------ .../relations/pooled_object_placer.py | 6 +- isaaclab_arena/relations/relation_solver.py | 22 +-- .../relations/relation_solver_state.py | 28 ++-- isaaclab_arena/relations/relations.py | 20 +-- .../{assets => tests}/dummy_embodiment.py | 19 ++- isaaclab_arena/tests/test_placement_events.py | 94 ++++++------ ...st_relation_solver_background_collision.py | 16 +- .../tests/test_relation_solver_embodiment.py | 20 ++- .../tests/test_relation_solver_interface.py | 46 +++--- .../placement_pool_ik_validation.py | 8 +- 27 files changed, 450 insertions(+), 519 deletions(-) rename isaaclab_arena/relations/{placement_entity.py => placement_asset.py} (64%) rename isaaclab_arena/{assets => tests}/dummy_embodiment.py (70%) 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 index ca65003b20..664a3dcf67 100644 --- a/isaaclab_arena/assets/dummy_object.py +++ b/isaaclab_arena/assets/dummy_object.py @@ -7,13 +7,13 @@ import torch import trimesh -from isaaclab_arena.relations.placement_entity import PlacementEntity +from isaaclab_arena.relations.placement_asset import PlacementAsset 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 -class DummyObject(PlacementEntity): +class DummyObject(PlacementAsset): """Dummy object for testing purposes without Isaac Sim dependencies.""" def __init__( @@ -36,16 +36,6 @@ 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) diff --git a/isaaclab_arena/assets/object.py b/isaaclab_arena/assets/object.py index fdeb669629..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 @@ -73,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 e59e1fd06a..9d3085cd59 100644 --- a/isaaclab_arena/assets/object_base.py +++ b/isaaclab_arena/assets/object_base.py @@ -7,13 +7,9 @@ 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 @@ -24,7 +20,7 @@ # 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.placement_entity import PlacementEntity +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.pose import Pose, PosePerEnv, PoseRange from isaaclab_arena.utils.velocity import Velocity @@ -36,7 +32,7 @@ ] -class ObjectBase(PlacementEntity, ABC): +class ObjectBase(PlacementAsset, ABC): """Parent class for (spawnable) Object and ObjectReference.""" def __init__( @@ -57,9 +53,6 @@ def __init__( self.object_cfg = None self.event_cfg = None - 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. @@ -91,14 +84,14 @@ def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: self.event_cfg = self._init_event_cfg() def set_placement_initial_pose(self, pose: Pose) -> None: - """Set a solved root pose without configuring an independent reset.""" + """Set the solved spawn pose without rebuilding the object reset event.""" self.initial_pose = pose if self.object_cfg is not None: self.object_cfg.init_state.pos = pose.position_xyz self.object_cfg.init_state.rot = pose.rotation_xyzw def has_pose_reset_event(self) -> bool: - """Return whether another reset event controls the root pose.""" + """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: diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index 74b03786e6..228a2b2dc0 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -6,23 +6,20 @@ from __future__ import annotations from collections.abc import Mapping -from typing import TYPE_CHECKING, Any +from typing import Any from isaaclab.envs import ManagerBasedRLMimicEnv from isaaclab.managers.recorder_manager import RecorderManagerBaseCfg from isaaclab_arena.embodiments.common.arm_mode import ArmMode -from isaaclab_arena.relations.placement_entity import PlacementEntity -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, quaternion_to_90_deg_z_quarters +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, PosePerEnv, PoseRange -if TYPE_CHECKING: - import trimesh - -class EmbodimentBase(PlacementEntity): +class EmbodimentBase(PlacementAsset): name: str | None = None tags: list[str] = ["embodiment"] @@ -57,27 +54,19 @@ def __init__( self.termination_cfg: Any | None = None def get_bounding_box(self) -> AxisAlignedBoundingBox: - """Return root-relative bounds for the USD-authored articulation pose.""" + """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" - spawn = self.scene_config.robot.spawn + 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: Compute bounds at configured initial joint positions when joint_pos is not None. + # 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_world_bounding_box(self) -> AxisAlignedBoundingBox: - """Return bounds transformed by the configured root pose.""" - bounding_box = self.get_bounding_box() - if self.initial_pose is None: - return bounding_box - quarters = quaternion_to_90_deg_z_quarters(self.initial_pose.rotation_xyzw) - return bounding_box.rotated_90_around_z(quarters).translated(self.initial_pose.position_xyz) - - def get_collision_mesh(self) -> trimesh.Trimesh | None: - """Return no mesh because embodiment placement uses bounds.""" - def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: """Set the embodiment root pose.""" assert isinstance(pose, Pose), "Embodiments require one root Pose" @@ -89,7 +78,7 @@ def supports_per_env_initial_pose(self) -> bool: def set_joint_initial_pos(self, joint_pos: Mapping[str, float]) -> None: """Update the robot's initial joint positions by joint name.""" - assert self.scene_config is not None, "scene_config.robot must be populated before setting joint positions" + 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) @@ -173,7 +162,7 @@ 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: - assert scene_config is not None, "scene_config.robot must be populated before setting the root pose" + 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 @@ -189,6 +178,10 @@ def get_termination_cfg(self) -> Any: def get_embodiment_name_in_scene(self) -> str: return "robot" + def get_scene_name(self) -> str: + """Return the embodiment's Isaac Lab scene key.""" + return self.get_embodiment_name_in_scene() + def get_ee_frame_name(self, arm_mode: ArmMode) -> str: # In case of multiple ee frames one can use self.mimic_arm_mode to get the correct ee frame name return "" diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py index bbaa7e62c6..a3add11dad 100644 --- a/isaaclab_arena/environments/arena_env_builder.py +++ b/isaaclab_arena/environments/arena_env_builder.py @@ -71,7 +71,7 @@ def _solve_relations(self) -> None: """Solve spatial relations for scene objects and the embodiment. This method: - 1. Collects placement entities that have relations + 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 @@ -82,15 +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. """ - placement_entities = list(self.arena_env.scene.get_objects_with_relations()) - scene_entity_names = {entity.name: entity.name for entity in placement_entities} + 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_entities.append(embodiment) - scene_entity_names[embodiment.name] = embodiment.get_embodiment_name_in_scene() + placement_assets.append(embodiment) placer_params = self.arena_env.placer_params if placer_params is None: @@ -106,11 +105,10 @@ 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( - placement_entities, + placement_assets, num_envs=self.cfg.num_envs, placer_params=placer_params, scene_assets=self.arena_env.scene.assets.values(), - scene_entity_names=scene_entity_names, ) def get_all_variations(self) -> dict[str, list[VariationBase]]: diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py index 5e2b6e371e..c4ed90350f 100644 --- a/isaaclab_arena/environments/relation_solver_interface.py +++ b/isaaclab_arena/environments/relation_solver_interface.py @@ -6,14 +6,14 @@ from __future__ import annotations import copy -from collections.abc import Iterable, Mapping +from collections.abc import Iterable from typing import TYPE_CHECKING 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_pose_from_layout, - place_entities_from_layouts, + place_assets_from_layouts, solve_and_place_objects, ) from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer @@ -26,7 +26,7 @@ from isaaclab_arena.assets.asset import Asset from isaaclab_arena.assets.object_set import RigidObjectSet from isaaclab_arena.relations.collision_object import CollisionObject - from isaaclab_arena.relations.placement_entity import PlacementEntity + from isaaclab_arena.relations.placement_asset import PlacementAsset from isaaclab_arena.relations.placement_result import PlacementResult @@ -41,17 +41,16 @@ def _get_passive_collision_objects( def solve_and_apply_relation_placement( - objects: list[PlacementEntity], + assets: list[PlacementAsset], num_envs: int, placer_params: ObjectPlacerParams | None = None, collision_objects: list[CollisionObject] | None = None, scene_assets: Iterable[Asset | RigidObjectSet] | None = None, - scene_entity_names: Mapping[str, str] | 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: Entities 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,21 +58,19 @@ def solve_and_apply_relation_placement( or relation-constrained. scene_assets: Optional scene assets to scan for passive collision objects when collision_objects is not supplied. - scene_entity_names: Isaac Lab scene name for each placement entity. Returns: 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 - entity_names = {obj.name for obj in objects} - assert len(entity_names) == len(objects), "Placement entity names must be unique" - if scene_entity_names is None: - scene_entity_names = {obj.name: obj.name for obj in objects} - assert set(scene_entity_names) == entity_names, "scene_entity_names must contain every placement entity name" + 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() @@ -89,11 +86,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 ), ) 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, @@ -110,25 +107,24 @@ 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, - scene_entity_names=scene_entity_names, ) def _should_include_background_mesh( - objects: list[PlacementEntity], + 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 @@ -137,129 +133,125 @@ def _should_include_background_mesh( def _apply_relation_placement_result( - objects: list[PlacementEntity], + assets: list[PlacementAsset], placer_params: ObjectPlacerParams, placement_pool: PooledObjectPlacer, num_envs: int, - scene_entity_names: Mapping[str, str], ) -> 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, - scene_entity_names=scene_entity_names, + anchor_assets=anchor_assets, ) - if all(obj.supports_per_env_initial_pose() for obj in objects): + # Objects can store PosePerEnv, so their reset events restore fixed per-env + # poses. Scenes containing any asset without per-env pose support use one + # coordinated event to restore layouts[env_id]. + if all(asset.supports_per_env_initial_pose() for asset in assets): _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, ) return None return _apply_static_spawn_pose( - objects=objects, + assets=assets, placement_pool=placement_pool, - anchor_objects_set=anchor_objects_set, + anchor_assets=anchor_assets, num_envs=num_envs, - scene_entity_names=scene_entity_names, ) def _apply_dynamic_spawn_pose( - objects: list[PlacementEntity], + assets: list[PlacementAsset], placement_pool: PooledObjectPlacer, - anchor_objects_set: set[PlacementEntity], - scene_entity_names: Mapping[str, str], + 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] - _set_placement_initial_poses(objects, anchor_objects_set, layout) + _set_placement_initial_poses(assets, anchor_assets, layout) return EventTermCfg( func=solve_and_place_objects, mode="reset", params={ - "objects": objects, + "assets": assets, "placement_pool": placement_pool, - "scene_entity_names": scene_entity_names, }, ) def _apply_static_spawn_pose( - objects: list[PlacementEntity], + assets: list[PlacementAsset], placement_pool: PooledObjectPlacer, - anchor_objects_set: set[PlacementEntity], + anchor_assets: set[PlacementAsset], num_envs: int, - scene_entity_names: Mapping[str, str], ) -> EventTermCfg: - """Return a reset event that restores one fixed layout per environment.""" + """Return a coordinated reset event that restores one fixed layout per environment.""" from isaaclab.managers import EventTermCfg layouts = placement_pool.sample_with_replacement(num_envs) - _set_placement_initial_poses(objects, anchor_objects_set, layouts[0]) + _set_placement_initial_poses(assets, anchor_assets, layouts[0]) return EventTermCfg( - func=place_entities_from_layouts, + func=place_assets_from_layouts, mode="reset", params={ - "objects": objects, + "assets": assets, "layouts": layouts, - "scene_entity_names": scene_entity_names, }, ) def _set_placement_initial_poses( - objects: list[PlacementEntity], - anchor_objects_set: set[PlacementEntity], + assets: list[PlacementAsset], + anchor_assets: set[PlacementAsset], layout: PlacementResult, ) -> None: - """Seed non-anchor entities from one placement layout.""" - for obj in objects: - if obj in anchor_objects_set: + """Seed the spawn pose while preserving coordinated reset ownership.""" + for asset in assets: + if asset in anchor_assets: continue - pose = get_pose_from_layout(obj, layout) - obj.set_placement_initial_pose(pose) + pose = get_pose_from_layout(asset, layout) + asset.set_placement_initial_pose(pose) def _apply_static_initial_poses( - objects: list[PlacementEntity], + assets: list[PlacementAsset], placement_pool: PooledObjectPlacer, - anchor_objects_set: set[PlacementEntity], + 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 - poses = [get_pose_from_layout(obj, layouts[env_idx]) for env_idx in range(num_envs)] - 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[PlacementEntity], - anchor_objects_set: set[PlacementEntity], + 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 obj.has_pose_reset_event()), ( - 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 e3f92670cb..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 placement entities. +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.relations.placement_entity import PlacementEntity + from isaaclab_arena.relations.placement_asset import PlacementAsset -def has_heterogeneous_objects(objects: list[PlacementEntity]) -> 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[PlacementEntity], 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[PlacementEntity], num_envs: int, plac variant_set_idx += 1 -def get_bounding_box_per_env(obj: PlacementEntity, 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[PlacementEntity, 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[PlacementEntity, 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[PlacementEntity, Ax for obj, bbox in self.object_bboxes.items() } - def get_bounding_boxes_for_all_envs(self) -> list[dict[PlacementEntity, 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[PlacementEntity, AxisAlig def get_bounding_boxes_for_solver_candidates( self, candidates_per_env: int - ) -> dict[PlacementEntity, 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[PlacementEntity], 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 bf7a06ac4f..1adf8fac13 100644 --- a/isaaclab_arena/relations/mesh_pair_cache.py +++ b/isaaclab_arena/relations/mesh_pair_cache.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject - from isaaclab_arena.relations.placement_entity import PlacementEntity + 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: PlacementEntity + subject: PlacementAsset """Subject (sphere source) object.""" - obstacle: PlacementEntity | 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[PlacementEntity] + pair_subject_objs: list[PlacementAsset] """(P,) subject (sphere source) object reference per pair.""" - pair_obstacle_objs: list[PlacementEntity | 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 37877b1bad..4b9d4e6fb3 100644 --- a/isaaclab_arena/relations/no_overlap_aabb.py +++ b/isaaclab_arena/relations/no_overlap_aabb.py @@ -18,7 +18,7 @@ if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject - from isaaclab_arena.relations.placement_entity import PlacementEntity + 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[PlacementEntity | 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: PlacementEntity, + 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: PlacementEntity, - b: PlacementEntity, + 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: PlacementEntity, + 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 f9728be559..e5fb52b8f6 100644 --- a/isaaclab_arena/relations/no_overlap_mesh.py +++ b/isaaclab_arena/relations/no_overlap_mesh.py @@ -23,7 +23,7 @@ if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject - from isaaclab_arena.relations.placement_entity import PlacementEntity + 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[PlacementEntity, 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[PlacementEntity | 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: PlacementEntity, + 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 6e0275d039..baa8e760eb 100644 --- a/isaaclab_arena/relations/object_placer.py +++ b/isaaclab_arena/relations/object_placer.py @@ -30,7 +30,7 @@ if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject - from isaaclab_arena.relations.placement_entity import PlacementEntity + 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[PlacementEntity, 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[PlacementEntity, 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[PlacementEntity], + 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[PlacementEntity], + 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[PlacementEntity], - ) -> tuple[set[PlacementEntity], 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[PlacementEntity], - anchor_objects_set: set[PlacementEntity], + 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[PlacementEntity, tuple[float, float, float]]] = [] - orientations_per_candidate: list[dict[PlacementEntity, 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[PlacementEntity], - anchor_objects: set[PlacementEntity], - env_bboxes: dict[PlacementEntity, AxisAlignedBoundingBox], + objects: list[PlacementAsset], + anchor_objects: set[PlacementAsset], + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], generator: torch.Generator | None = None, - ) -> dict[PlacementEntity, 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[PlacementEntity, 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: PlacementEntity, - env_bboxes: dict[PlacementEntity, 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[PlacementEntity], - anchor_objects: set[PlacementEntity], + objects: list[PlacementAsset], + anchor_objects: set[PlacementAsset], generator: torch.Generator | None = None, - ) -> dict[PlacementEntity, 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[PlacementEntity, 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[PlacementEntity, tuple[float, float, float]]], - orientations_per_candidate: list[dict[PlacementEntity, 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[PlacementEntity], - candidate_bboxes: dict[PlacementEntity, AxisAlignedBoundingBox], - orientations_per_candidate: list[dict[PlacementEntity, float]], - ) -> dict[PlacementEntity, 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[PlacementEntity, 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[PlacementEntity, AxisAlignedBoundingBox], + bboxes: dict[PlacementAsset, AxisAlignedBoundingBox], candidate_idx: int, - ) -> dict[PlacementEntity, 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: PlacementEntity, - anchor_objects: set[PlacementEntity], + parent: PlacementAsset, + anchor_objects: set[PlacementAsset], anchor_bbox: AxisAlignedBoundingBox, - env_bboxes: dict[PlacementEntity, 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: PlacementEntity, - anchor_objects: set[PlacementEntity], + obj: PlacementAsset, + anchor_objects: set[PlacementAsset], anchor_bbox: AxisAlignedBoundingBox, - env_bboxes: dict[PlacementEntity, 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[PlacementEntity, tuple[float, float, float]]], - orientations: list[dict[PlacementEntity, float]], - bboxes: list[dict[PlacementEntity, 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. @@ -700,9 +700,9 @@ def _passes_required_checks( def _apply_poses( self, - positions_per_env: list[dict[PlacementEntity, tuple[float, float, float]]], - anchor_objects: set[PlacementEntity], - orientations_per_env: list[dict[PlacementEntity, 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_entity.py b/isaaclab_arena/relations/placement_asset.py similarity index 64% rename from isaaclab_arena/relations/placement_entity.py rename to isaaclab_arena/relations/placement_asset.py index 830dbf39c3..2361138362 100644 --- a/isaaclab_arena/relations/placement_entity.py +++ b/isaaclab_arena/relations/placement_asset.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Shared model for entities whose poses are relation-solved.""" +"""Shared model for assets whose poses are relation-solved.""" from __future__ import annotations @@ -13,6 +13,7 @@ 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: @@ -21,7 +22,7 @@ from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox -class PlacementEntity(Asset, ABC): +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: @@ -30,15 +31,15 @@ def __init__(self, name: str, tags: list[str] | None = None, **kwargs) -> None: self.relations: list[RelationBase] = [] # None delegates collision-mode selection to the solver. self.collision_mode: CollisionMode | None = None - # Mesh collision uses a convex hull when source geometry is not watertight. + # 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 entity.""" + """Attach a relation to the asset.""" self.relations.append(relation) def get_relations(self) -> list[RelationBase]: - """Return all relations attached to the entity.""" + """Return all relations attached to the asset.""" return self.relations def get_spatial_relations(self) -> list[RelationBase]: @@ -47,7 +48,7 @@ def get_spatial_relations(self) -> list[RelationBase]: @property def is_anchor(self) -> bool: - """Return whether the entity is fixed during relation solving.""" + """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: @@ -55,15 +56,19 @@ def get_initial_pose(self) -> Pose | PoseRange | PosePerEnv | None: return self.initial_pose def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: - """Set the configured root pose.""" + """Set the configured root pose. + + ``PoseRange`` and ``PosePerEnv`` support is subclass-specific. Callers + assigning ``PosePerEnv`` must first check ``supports_per_env_initial_pose()``. + """ self.initial_pose = pose def set_placement_initial_pose(self, pose: Pose) -> None: - """Set a solved root pose without changing reset ownership.""" + """Set the solved spawn pose without changing reset ownership.""" self.set_initial_pose(pose) def has_pose_reset_event(self) -> bool: - """Return whether the entity owns a root-pose reset event.""" + """Return whether the asset owns a root-pose reset event.""" return False def supports_per_env_initial_pose(self) -> bool: @@ -74,10 +79,17 @@ def supports_per_env_initial_pose(self) -> bool: def get_bounding_box(self) -> AxisAlignedBoundingBox: """Return root-relative axis-aligned bounds.""" - @abstractmethod def get_world_bounding_box(self) -> AxisAlignedBoundingBox: - """Return axis-aligned bounds transformed by the root pose.""" + """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) - @abstractmethod def get_collision_mesh(self) -> trimesh.Trimesh | None: - """Return a root-relative collision mesh when available.""" + """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 e8bb189fe0..79e619fdb0 100644 --- a/isaaclab_arena/relations/placement_events.py +++ b/isaaclab_arena/relations/placement_events.py @@ -6,7 +6,6 @@ from __future__ import annotations import torch -from collections.abc import Mapping from typing import TYPE_CHECKING from isaaclab_arena.relations.relations import RotateAroundSolution, get_anchor_objects @@ -17,7 +16,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv - from isaaclab_arena.relations.placement_entity import PlacementEntity + 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 @@ -43,94 +42,78 @@ def get_placement_pool(env) -> PooledObjectPlacer | None: return term_cfg.params.get("placement_pool") -def get_placement_scene_entity_names(env) -> Mapping[str, str] | None: - """Return the scene-name map stored on the placement reset event.""" - try: - term_cfg = env.unwrapped.event_manager.get_term_cfg(PLACEMENT_RESET_EVENT_NAME) - except ValueError: - return None - return term_cfg.params.get("scene_entity_names") - - -def get_rotation_xyzw(obj: PlacementEntity) -> 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[PlacementEntity], -) -> dict[PlacementEntity, 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_pose_from_layout(obj: PlacementEntity, layout: PlacementResult) -> Pose: - """Return an entity pose from a solved layout.""" - assert obj in layout.positions, f"Placement layout is missing non-anchor entity '{obj.name}'" - base_rotation = get_rotation_xyzw(obj) +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(obj, marker_yaw) + 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[obj], rotation_xyzw=rotation) + return Pose(position_xyz=layout.positions[asset], rotation_xyzw=rotation) -def get_movable_object_names( - objects: list[PlacementEntity], - anchor_objects_set: set[PlacementEntity], - scene_entity_names: Mapping[str, str], +def get_movable_asset_names( + assets: list[PlacementAsset], + anchor_assets: set[PlacementAsset], ) -> list[str]: - """Return scene names for non-anchor placement entities.""" - return [scene_entity_names[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_layout_to_sim( env: ManagerBasedEnv, env_id: int, result: PlacementResult, - anchor_objects_set: set[PlacementEntity], - base_rotations: dict[PlacementEntity, tuple[float, float, float, float]], - scene_entity_names: Mapping[str, str] | None = None, + 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. - scene_entity_names: Isaac Lab scene name for each placement entity. Entity - names are used directly when omitted. + anchor_assets: The set of anchor assets. + base_rotations: The base rotations for all assets. """ - if scene_entity_names is None: - scene_entity_names = {obj.name: obj.name for obj in base_rotations} env_id_tensor = torch.tensor([env_id], device=env.device) zero_velocity = Velocity.zero().to_tensor(device=env.device).unsqueeze(0) - missing_entities = [ - obj.name for obj in base_rotations if obj not in anchor_objects_set and obj not in result.positions + missing_assets = [ + asset.name for asset in base_rotations if asset not in anchor_assets and asset not in result.positions ] - assert not missing_entities, f"Placement layout is missing non-anchor entities: {missing_entities}" - for obj in result.positions: - if obj in anchor_objects_set: + 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[scene_entity_names[obj.name]] - pose = get_pose_from_layout(obj, result) + scene_asset = env.scene[asset.get_scene_name()] + pose = get_pose_from_layout(asset, result) pose_tensor = pose.to_tensor(device=env.device).unsqueeze(0) pose_tensor[0, :3] += env.scene.env_origins[env_id, :] - asset.write_root_pose_to_sim(pose_tensor, env_ids=env_id_tensor) - asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_id_tensor) + scene_asset.write_root_pose_to_sim(pose_tensor, env_ids=env_id_tensor) + scene_asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_id_tensor) def solve_and_place_objects( env: ManagerBasedEnv, env_ids: torch.Tensor | None, - objects: list[PlacementEntity], + assets: list[PlacementAsset], placement_pool: PooledObjectPlacer, - scene_entity_names: Mapping[str, str], ) -> None: """Coordinated reset event that draws layouts from the pool and writes poses. @@ -141,9 +124,8 @@ 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. - scene_entity_names: Isaac Lab scene name for each placement entity. """ if env_ids is None or len(env_ids) == 0: return @@ -153,8 +135,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] @@ -163,37 +145,34 @@ 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, scene_entity_names) + # Only write non-anchor assets to the sim. + write_layout_to_sim(env, cur_env, result, anchor_assets, base_rotations) -def place_entities_from_layouts( +def place_assets_from_layouts( env: ManagerBasedEnv, env_ids: torch.Tensor | None, - objects: list[PlacementEntity], + assets: list[PlacementAsset], layouts: list[PlacementResult], - scene_entity_names: Mapping[str, str], ) -> None: """Restore one fixed placement layout per environment. Args: env: The Isaac Lab environment. env_ids: Environment indices to restore. - objects: Entities participating in relation solving. + assets: Assets participating in relation solving. layouts: Fixed layout indexed by environment. - scene_entity_names: Isaac Lab scene name for each placement entity. """ if env_ids is None or len(env_ids) == 0: return assert len(layouts) == env.scene.env_origins.shape[0], "Static layouts must match the scene environment count" - 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 env_id in env_ids.tolist(): write_layout_to_sim( env, env_id, layouts[env_id], - anchor_objects_set, + anchor_assets, base_rotations, - scene_entity_names, ) diff --git a/isaaclab_arena/relations/placement_pool_validation.py b/isaaclab_arena/relations/placement_pool_validation.py index 1e437e2330..98bcb4fb9f 100644 --- a/isaaclab_arena/relations/placement_pool_validation.py +++ b/isaaclab_arena/relations/placement_pool_validation.py @@ -5,15 +5,13 @@ from __future__ import annotations -from collections.abc import Mapping from typing import TYPE_CHECKING 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, - get_placement_scene_entity_names, write_layout_to_sim, ) from isaaclab_arena.relations.placement_validation import PlacementCheck @@ -23,7 +21,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv - from isaaclab_arena.relations.placement_entity import PlacementEntity + 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 @@ -34,9 +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[PlacementEntity, tuple[float, float, float, float]], - scene_entity_names: Mapping[str, str], + 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. @@ -52,9 +49,8 @@ def _write_layout_to_envs_for_episode_index( env.unwrapped, env_id, layout, - anchor_objects_set, + anchor_assets, base_rotations, - scene_entity_names, ) layouts_written.append((env_id, layout)) return layouts_written @@ -89,7 +85,6 @@ def validate_pool_layouts( placement_pool: PooledObjectPlacer | None = None, settle_params: PhysicsSettleParams | None = None, render: bool = False, - scene_entity_names: Mapping[str, str] | None = None, ) -> list[tuple[int, int, PlacementValidationResults]] | None: """Physics-validate every layout in a placement pool, recording the result on its validation results. @@ -103,13 +98,11 @@ def validate_pool_layouts( settle_params: Settle-check tuning params. Defaults to ``PhysicsSettleParams()`` when omitted. render: When True, render each settle step so the sweep is visible in the GUI. Defaults to False. - scene_entity_names: Isaac Lab scene name for each placement entity. Returns: ``(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. """ - placement_pool_from_event = placement_pool is None if placement_pool is None: placement_pool = get_placement_pool(env) if placement_pool is None: @@ -117,15 +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) - if scene_entity_names is None: - scene_entity_names = get_placement_scene_entity_names(env) - if scene_entity_names is None: - assert not placement_pool_from_event, "Placement reset event is missing scene_entity_names" - scene_entity_names = {obj.name: obj.name for obj in objects} - movable_object_names = get_movable_object_names(objects, anchor_objects_set, scene_entity_names) + 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() @@ -148,9 +136,8 @@ def validate_pool_layouts( layouts_per_env, num_envs, episode_index, - anchor_objects_set, + anchor_assets, base_rotations, - scene_entity_names, ) 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 be9a56bc48..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.relations.placement_entity import PlacementEntity + from isaaclab_arena.relations.placement_asset import PlacementAsset from isaaclab_arena.relations.placement_validation import PlacementValidationResults @dataclass class PlacementResult: - """Solved entity layout for one environment.""" + """Solved asset layout for one environment.""" validation_results: PlacementValidationResults """Validation checklist for the placement.""" - positions: dict[PlacementEntity, tuple[float, float, float]] - """Final positions for each entity.""" + 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[PlacementEntity, float] = field(default_factory=dict) - """Placement-computed world Z-yaws. Omitted entities 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 522bded01b..2b3e0cfbef 100644 --- a/isaaclab_arena/relations/pooled_object_placer.py +++ b/isaaclab_arena/relations/pooled_object_placer.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject - from isaaclab_arena.relations.placement_entity import PlacementEntity + from isaaclab_arena.relations.placement_asset import PlacementAsset @dataclass @@ -70,7 +70,7 @@ class PooledObjectPlacer: def __init__( self, - objects: list[PlacementEntity], + 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[PlacementEntity]: + 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 0ffc4cc3b9..f1a3a069b7 100644 --- a/isaaclab_arena/relations/relation_solver.py +++ b/isaaclab_arena/relations/relation_solver.py @@ -24,7 +24,7 @@ if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject from isaaclab_arena.relations.mesh_pair_cache import MeshPairCache - from isaaclab_arena.relations.placement_entity import PlacementEntity + 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[PlacementEntity, 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[PlacementEntity], - initial_positions: list[dict[PlacementEntity, tuple[float, float, float]]], - env_bboxes: dict[PlacementEntity, 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[PlacementEntity, float]] | None = None, + orientations: list[dict[PlacementAsset, float]] | None = None, collision_objects: list[CollisionObject] | None = None, - ) -> list[dict[PlacementEntity, tuple[float, float, float]]]: + ) -> list[dict[PlacementAsset, tuple[float, float, float]]]: """Solve for optimal positions of all objects. Args: - objects: Placement entities to solve. Must include at least one entity + 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[PlacementEntity]) -> 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[PlacementEntity]) -> None: def _print_relation_debug( - obj: PlacementEntity, + 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: PlacementEntity, + 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 134b031336..e18ae5a160 100644 --- a/isaaclab_arena/relations/relation_solver_state.py +++ b/isaaclab_arena/relations/relation_solver_state.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: from isaaclab_arena.relations.collision_object import CollisionObject - from isaaclab_arena.relations.placement_entity import PlacementEntity + from isaaclab_arena.relations.placement_asset import PlacementAsset class RelationSolverState: @@ -28,16 +28,16 @@ class RelationSolverState: def __init__( self, - objects: list[PlacementEntity], - initial_positions: list[dict[PlacementEntity, tuple[float, float, float]]], + objects: list[PlacementAsset], + initial_positions: list[dict[PlacementAsset, tuple[float, float, float]]], device: torch.device | None = None, - env_bboxes: dict[PlacementEntity, AxisAlignedBoundingBox] | None = None, + env_bboxes: dict[PlacementAsset, AxisAlignedBoundingBox] | None = None, collision_objects: list[CollisionObject] | None = None, ): """Initialize optimization state. Args: - objects: Placement entities 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[PlacementEntity] = 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[PlacementEntity, 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[PlacementEntity | 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[PlacementEntity]: + def optimizable_objects(self) -> list[PlacementAsset]: """List of optimizable objects (excludes anchors).""" return self._optimizable_objects @property - def anchor_objects(self) -> set[PlacementEntity]: + 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: PlacementEntity) -> 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: PlacementEntity) -> torch.Tensor: opt_idx = self._global_to_opt_idx[idx] return self._optimizable_positions[:, opt_idx, :] - def get_fixed_obstacle_world_bbox(self, obj: PlacementEntity | 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: PlacementEntity) -> 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[PlacementEntity, 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 045e73a60b..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.relations.placement_entity import PlacementEntity + 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: PlacementEntity, objects: set[PlacementEntity]) -> 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: PlacementEntity, 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: PlacementEntity): + 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: PlacementEntity, objects: set[PlacementEntity]) -> 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: PlacementEntity, + 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: PlacementEntity, + 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: PlacementEntity, + 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[PlacementEntity]) -> list[PlacementEntity]: +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[PlacementEntity]) -> list[PlacementEntity]: 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/assets/dummy_embodiment.py b/isaaclab_arena/tests/dummy_embodiment.py similarity index 70% rename from isaaclab_arena/assets/dummy_embodiment.py rename to isaaclab_arena/tests/dummy_embodiment.py index 651ba0bc66..55257a19fc 100644 --- a/isaaclab_arena/assets/dummy_embodiment.py +++ b/isaaclab_arena/tests/dummy_embodiment.py @@ -9,12 +9,12 @@ import trimesh -from isaaclab_arena.relations.placement_entity import PlacementEntity -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, quaternion_to_90_deg_z_quarters +from isaaclab_arena.relations.placement_asset import PlacementAsset +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose -class DummyEmbodiment(PlacementEntity): +class DummyEmbodiment(PlacementAsset): """Embodiment geometry without simulator dependencies.""" def __init__( @@ -23,23 +23,18 @@ def __init__( 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 def get_bounding_box(self) -> AxisAlignedBoundingBox: """Return root-relative bounds.""" return self.bounding_box - def get_world_bounding_box(self) -> AxisAlignedBoundingBox: - """Return bounds transformed by the root pose.""" - 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_collision_mesh(self) -> trimesh.Trimesh | None: """Return the configured collision mesh.""" return self.collision_mesh @@ -47,3 +42,7 @@ def get_collision_mesh(self) -> trimesh.Trimesh | None: def supports_per_env_initial_pose(self) -> bool: """Return False because the dummy stores one root pose.""" return False + + def get_scene_name(self) -> str: + """Return the configured scene key.""" + return self.scene_name diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index 1d39f09d4d..067f5cb76c 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -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 @@ -162,15 +156,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, - scene_entity_names={obj.name: obj.name for obj in objects}, ) 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 @@ -200,11 +191,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: @@ -215,7 +212,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, ) @@ -224,24 +221,25 @@ 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(), - scene_entity_names={"desk": "desk", "box1": "box1"}, ) 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 test_static_layout_event_writes_embodiment_to_mapped_scene_asset(): - from isaaclab_arena.assets.dummy_embodiment import DummyEmbodiment - from isaaclab_arena.relations.placement_events import place_entities_from_layouts +def test_static_layout_event_writes_embodiment_to_configured_scene_asset(): + from isaaclab_arena.relations.placement_events import place_assets_from_layouts from isaaclab_arena.relations.placement_result import PlacementResult + from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox 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", ) layouts = [ PlacementResult( @@ -259,12 +257,11 @@ def test_static_layout_event_writes_embodiment_to_mapped_scene_asset(): ] env = _make_mock_env(num_envs=2) - place_entities_from_layouts( + place_assets_from_layouts( env, torch.tensor([1]), - objects=[robot], + assets=[robot], layouts=layouts, - scene_entity_names={"droid": "robot"}, ) assert "droid" not in env._assets @@ -272,17 +269,32 @@ def test_static_layout_event_writes_embodiment_to_mapped_scene_asset(): assert torch.allclose(pose[0, :3], torch.tensor([0.4, 0.5, 0.0])) with pytest.raises(AssertionError, match="Static layouts must match"): - place_entities_from_layouts( + place_assets_from_layouts( env, torch.tensor([0]), - objects=[robot], + assets=[robot], layouts=layouts[:1], - scene_entity_names={"droid": "robot"}, ) +def test_static_layout_event_rejects_missing_non_anchor_assets(): + from isaaclab_arena.relations.placement_events import place_assets_from_layouts + 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"): + place_assets_from_layouts(env, torch.tensor([0]), assets=[box1], layouts=[layout]) + + 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: @@ -332,8 +344,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 @@ -351,8 +361,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 @@ -370,8 +378,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 @@ -455,14 +461,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 @@ -493,8 +498,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] @@ -515,8 +520,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 @@ -538,8 +541,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 @@ -557,8 +558,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 @@ -657,10 +656,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}, - scene_entity_names={"desk": "desk", "box": "box"}, + anchor_assets={anchor}, ) assert pool.sample_count == 1 @@ -708,9 +706,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, ) diff --git a/isaaclab_arena/tests/test_relation_solver_background_collision.py b/isaaclab_arena/tests/test_relation_solver_background_collision.py index f27a59b778..36a5c4a636 100644 --- a/isaaclab_arena/tests/test_relation_solver_background_collision.py +++ b/isaaclab_arena/tests/test_relation_solver_background_collision.py @@ -42,11 +42,7 @@ 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. - """ + """Return an obstacle narrower than the desk to preserve a non-zero escape gradient.""" from isaaclab_arena.assets.dummy_object import DummyObject from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -553,14 +549,12 @@ def fake_solve_and_apply_relation_placement( placer_params, collision_objects=None, scene_assets=None, - scene_entity_names=None, ): calls["objects"] = objects calls["num_envs"] = num_envs calls["placer_params"] = placer_params calls["scene_assets"] = list(scene_assets) calls["collision_objects"] = collision_objects - calls["scene_entity_names"] = scene_entity_names return "placement_event" monkeypatch.setattr(builder_module, "solve_and_apply_relation_placement", fake_solve_and_apply_relation_placement) @@ -575,7 +569,6 @@ def fake_solve_and_apply_relation_placement( assert calls["placer_params"] is placer_params assert calls["scene_assets"] == [background_collision] assert calls["collision_objects"] is None - assert calls["scene_entity_names"]["object"] == "object" assert builder._placement_event_cfg == "placement_event" @@ -601,12 +594,10 @@ def fake_solve_and_apply_relation_placement( placer_params, collision_objects=None, scene_assets=None, - scene_entity_names=None, ): calls["objects"] = objects calls["scene_assets"] = list(scene_assets) calls["collision_objects"] = collision_objects - calls["scene_entity_names"] = scene_entity_names monkeypatch.setattr(builder_module, "solve_and_apply_relation_placement", fake_solve_and_apply_relation_placement) arena_env = SimpleNamespace(scene=Scene(), embodiment=None, placer_params=None) @@ -617,7 +608,6 @@ def fake_solve_and_apply_relation_placement( assert calls["objects"] == [] assert calls["scene_assets"] == [] assert calls["collision_objects"] is None - assert calls["scene_entity_names"] == {} def test_arena_env_builder_includes_embodiment_relations(monkeypatch): @@ -641,9 +631,6 @@ class Embodiment: def get_relations(self): return [object()] - def get_embodiment_name_in_scene(self): - return "robot" - def fake_solve_and_apply_relation_placement(*args, **kwargs): calls.update(kwargs) calls["objects"] = args[0] @@ -655,7 +642,6 @@ def fake_solve_and_apply_relation_placement(*args, **kwargs): ArenaEnvBuilder(arena_env, ArenaEnvBuilderCfg())._solve_relations() assert calls["objects"] == [embodiment] - assert calls["scene_entity_names"] == {"droid": "robot"} def test_relation_placement_includes_background_mesh_for_object_mesh_override(monkeypatch): diff --git a/isaaclab_arena/tests/test_relation_solver_embodiment.py b/isaaclab_arena/tests/test_relation_solver_embodiment.py index 0b53debeaf..02a3c0a246 100644 --- a/isaaclab_arena/tests/test_relation_solver_embodiment.py +++ b/isaaclab_arena/tests/test_relation_solver_embodiment.py @@ -5,13 +5,15 @@ """Relation placement tests for embodiments.""" +import torch + import pytest -from isaaclab_arena.assets.dummy_embodiment import DummyEmbodiment 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.relations import IsAnchor, On +from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.pose import Pose @@ -52,3 +54,19 @@ def test_batched_embodiment_placement_requires_runtime_application(): with pytest.raises(AssertionError, match="cannot store per-environment poses"): ObjectPlacer(ObjectPlacerParams(placement_seed=3)).place([floor, robot], num_envs=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 1f5609f6b8..87e3c17da9 100644 --- a/isaaclab_arena/tests/test_relation_solver_interface.py +++ b/isaaclab_arena/tests/test_relation_solver_interface.py @@ -62,18 +62,26 @@ 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_entity_names(): +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_requires_complete_scene_name_map(): +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="must contain every placement entity"): - solve_and_apply_relation_placement([_make_desk()], num_envs=1, scene_entity_names={}) + 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(): @@ -121,16 +129,15 @@ def test_dynamic_spawn_pose_rejects_layout_missing_non_anchor(): box = _make_box() placement_pool = _FakePlacementPool([_fallback_layout(positions={})]) - with pytest.raises(AssertionError, match="missing non-anchor entity 'box'"): + with pytest.raises(AssertionError, match="missing non-anchor asset 'box'"): _apply_dynamic_spawn_pose( - objects=[desk, box], + assets=[desk, box], placement_pool=placement_pool, - anchor_objects_set={desk}, - scene_entity_names={"desk": "desk", "box": "box"}, + 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() @@ -138,21 +145,19 @@ def test_dynamic_spawn_pose_event_params_use_runtime_objects(): 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}, - scene_entity_names={"desk": "desk", "box": "box"}, + 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 - assert event_cfg.params["scene_entity_names"]["box"] == "box" def test_static_embodiment_placement_uses_coordinated_reset(): - from isaaclab_arena.assets.dummy_embodiment import DummyEmbodiment 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 Pose @@ -170,18 +175,17 @@ def test_static_embodiment_placement_uses_coordinated_reset(): ] event_cfg = _apply_relation_placement_result( - objects=[desk, robot], + assets=[desk, robot], placer_params=ObjectPlacerParams(resolve_on_reset=False), placement_pool=_FakePlacementPool(layouts), num_envs=2, - scene_entity_names={"desk": "desk", "robot": "robot"}, ) initial_pose = robot.get_initial_pose() assert isinstance(initial_pose, Pose) assert initial_pose.position_xyz == (0.1, 0.2, 0.0) assert event_cfg is not None - runtime_robot = event_cfg.params["objects"][1] + runtime_robot = event_cfg.params["assets"][1] assert runtime_robot in event_cfg.params["layouts"][0].positions assert event_cfg.params["layouts"][1].positions[runtime_robot] == (0.3, 0.4, 0.0) @@ -197,10 +201,10 @@ def test_static_initial_poses_reject_layout_missing_non_anchor(): _fallback_layout(positions={placed_box: (0.2, 0.0, 0.2)}), ]) - with pytest.raises(AssertionError, match="missing non-anchor entity 'missing_box'"): + with pytest.raises(AssertionError, match="missing non-anchor asset 'missing_box'"): _apply_static_initial_poses( - objects=[desk, missing_box, placed_box], + assets=[desk, missing_box, placed_box], placement_pool=placement_pool, - anchor_objects_set={desk}, + anchor_assets={desk}, num_envs=2, ) 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) From 1b9a32fcabb586cac81388f565d33036b0875fb5 Mon Sep 17 00:00:00 2001 From: zhx06 Date: Tue, 21 Jul 2026 20:09:35 -0700 Subject: [PATCH 03/19] rename functions of initialization Signed-off-by: zhx06 --- isaaclab_arena/assets/object_base.py | 16 +++++++++------- .../environments/relation_solver_interface.py | 18 +++++++++++------- isaaclab_arena/relations/placement_asset.py | 4 ++-- isaaclab_arena/tests/test_placement_events.py | 4 ++-- 4 files changed, 24 insertions(+), 18 deletions(-) diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py index 9d3085cd59..4ccd600b37 100644 --- a/isaaclab_arena/assets/object_base.py +++ b/isaaclab_arena/assets/object_base.py @@ -76,19 +76,21 @@ 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_placement_initial_pose(self, pose: Pose) -> None: - """Set the solved spawn pose without rebuilding the object reset event.""" - self.initial_pose = pose - if self.object_cfg is not None: - self.object_cfg.init_state.pos = pose.position_xyz - self.object_cfg.init_state.rot = pose.rotation_xyzw + 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.""" diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py index c4ed90350f..dd048a1cc0 100644 --- a/isaaclab_arena/environments/relation_solver_interface.py +++ b/isaaclab_arena/environments/relation_solver_interface.py @@ -181,9 +181,10 @@ def _apply_dynamic_spawn_pose( """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] - _set_placement_initial_poses(assets, anchor_assets, layout) + # 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, @@ -205,7 +206,10 @@ def _apply_static_spawn_pose( from isaaclab.managers import EventTermCfg layouts = placement_pool.sample_with_replacement(num_envs) - _set_placement_initial_poses(assets, anchor_assets, layouts[0]) + # Scene configs hold one pose per asset, so construction uses one valid layout. + # The reset event installs layouts[env_id] before the first environment step. + construction_layout = layouts[0] + _seed_spawn_config_from_layout(assets, anchor_assets, construction_layout) return EventTermCfg( func=place_assets_from_layouts, mode="reset", @@ -216,17 +220,17 @@ def _apply_static_spawn_pose( ) -def _set_placement_initial_poses( +def _seed_spawn_config_from_layout( assets: list[PlacementAsset], anchor_assets: set[PlacementAsset], layout: PlacementResult, ) -> None: - """Seed the spawn pose while preserving coordinated reset ownership.""" + """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_placement_initial_pose(pose) + asset.set_spawn_pose(pose) def _apply_static_initial_poses( diff --git a/isaaclab_arena/relations/placement_asset.py b/isaaclab_arena/relations/placement_asset.py index 2361138362..0d1c20b619 100644 --- a/isaaclab_arena/relations/placement_asset.py +++ b/isaaclab_arena/relations/placement_asset.py @@ -63,8 +63,8 @@ def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: """ self.initial_pose = pose - def set_placement_initial_pose(self, pose: Pose) -> None: - """Set the solved spawn pose without changing reset ownership.""" + def set_spawn_pose(self, pose: Pose) -> None: + """Set the root pose used when constructing the scene.""" self.set_initial_pose(pose) def has_pose_reset_event(self) -> bool: diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index 067f5cb76c..4d3eac3ec9 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -630,7 +630,7 @@ 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_placement_initial_pose(self, pose): + def set_spawn_pose(self, pose): self.object_cfg.init_state.pos = pose.position_xyz self.object_cfg.init_state.rot = pose.rotation_xyzw @@ -648,7 +648,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") From 5862b929f790e1e07306e0d2ccad427128ebb69e Mon Sep 17 00:00:00 2001 From: zhx06 Date: Thu, 23 Jul 2026 11:36:41 -0700 Subject: [PATCH 04/19] fix mesh mode for droid Signed-off-by: zhx06 --- isaaclab_arena/embodiments/embodiment_base.py | 29 ++++++++++- .../tests/test_embodiment_collision_mesh.py | 52 +++++++++++++++++++ 2 files changed, 80 insertions(+), 1 deletion(-) create mode 100644 isaaclab_arena/tests/test_embodiment_collision_mesh.py diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index 228a2b2dc0..b2aedf07fe 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -6,7 +6,7 @@ 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.recorder_manager import RecorderManagerBaseCfg @@ -18,6 +18,9 @@ from isaaclab_arena.utils.configclass import combine_configclass_instances from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange +if TYPE_CHECKING: + import trimesh + class EmbodimentBase(PlacementAsset): @@ -52,6 +55,8 @@ 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.""" def get_bounding_box(self) -> AxisAlignedBoundingBox: """Return root-relative bounds computed from the articulation's USD geometry.""" @@ -67,6 +72,28 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox: # 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 set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: """Set the embodiment root pose.""" assert isinstance(pose, Pose), "Embodiments require one root Pose" 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() From ae4b68285467fbd1ad54c9a50c32f2a493de5ebb Mon Sep 17 00:00:00 2001 From: zhx06 Date: Thu, 23 Jul 2026 17:25:02 -0700 Subject: [PATCH 05/19] morefactor and move utils Signed-off-by: zhx06 --- isaaclab_arena/embodiments/embodiment_base.py | 5 +- isaaclab_arena/tasks/lift_object_task.py | 8 +- .../{assets => tests}/dummy_object.py | 0 isaaclab_arena/tests/test_face_to.py | 2 +- .../tests/test_heterogeneous_placement.py | 2 +- isaaclab_arena/tests/test_mesh_collision.py | 2 +- .../tests/test_no_collision_loss.py | 2 +- .../tests/test_object_placer_init.py | 2 +- .../test_object_placer_reproducibility.py | 2 +- isaaclab_arena/tests/test_placement_events.py | 10 +-- isaaclab_arena/tests/test_position_limits.py | 2 +- .../tests/test_relation_loss_strategies.py | 2 +- .../test_relation_placement_from_yaml.py | 79 +++++++++++++++++++ ...st_relation_solver_background_collision.py | 14 ++-- .../tests/test_relation_solver_embodiment.py | 2 +- .../tests/test_relation_solver_interface.py | 4 +- .../tests/test_validate_placement.py | 2 +- isaaclab_arena_curobo/utils/planner_utils.py | 2 +- .../relations/dummy_object_placer_notebook.py | 36 ++++----- .../relations/example_object.py | 19 +++++ .../relation_solver_visualization_notebook.py | 14 ++-- 21 files changed, 153 insertions(+), 58 deletions(-) rename isaaclab_arena/{assets => tests}/dummy_object.py (100%) create mode 100644 isaaclab_arena/tests/test_relation_placement_from_yaml.py create mode 100644 isaaclab_arena_examples/relations/example_object.py diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index b2aedf07fe..bbf46b89fb 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -202,12 +202,9 @@ 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: - return "robot" - def get_scene_name(self) -> str: """Return the embodiment's Isaac Lab scene key.""" - return self.get_embodiment_name_in_scene() + return "robot" def get_ee_frame_name(self, arm_mode: ArmMode) -> str: # In case of multiple ee frames one can use self.mimic_arm_mode to get the correct ee frame name 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/assets/dummy_object.py b/isaaclab_arena/tests/dummy_object.py similarity index 100% rename from isaaclab_arena/assets/dummy_object.py rename to isaaclab_arena/tests/dummy_object.py 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 4d3eac3ec9..298f26a0ce 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 @@ -669,10 +669,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 @@ -722,11 +722,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 @@ -764,11 +764,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 @@ -800,11 +800,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 36a5c4a636..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( @@ -43,7 +43,7 @@ def _make_box(name: str = "box"): def _make_background(): """Return an obstacle narrower than the desk to preserve a non-zero escape gradient.""" - 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 @@ -59,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 @@ -647,11 +647,11 @@ def fake_solve_and_apply_relation_placement(*args, **kwargs): 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( @@ -690,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) @@ -736,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 index 02a3c0a246..9c64617161 100644 --- a/isaaclab_arena/tests/test_relation_solver_embodiment.py +++ b/isaaclab_arena/tests/test_relation_solver_embodiment.py @@ -9,11 +9,11 @@ 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.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 diff --git a/isaaclab_arena/tests/test_relation_solver_interface.py b/isaaclab_arena/tests/test_relation_solver_interface.py index 87e3c17da9..50b30e5592 100644 --- a/isaaclab_arena/tests/test_relation_solver_interface.py +++ b/isaaclab_arena/tests/test_relation_solver_interface.py @@ -9,8 +9,8 @@ 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 @@ -24,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( 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_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_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 From ce534d7b838034fcea0967b469a7ce7f5529ce46 Mon Sep 17 00:00:00 2001 From: zhx06 Date: Thu, 23 Jul 2026 20:39:42 -0700 Subject: [PATCH 06/19] fix rebase issues Signed-off-by: zhx06 --- isaaclab_arena/relations/object_placer.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py index baa8e760eb..d3a194df90 100644 --- a/isaaclab_arena/relations/object_placer.py +++ b/isaaclab_arena/relations/object_placer.py @@ -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]], From da384eef9a735003e3db0e9375523dbae95b751a Mon Sep 17 00:00:00 2001 From: zhx06 Date: Thu, 23 Jul 2026 21:35:28 -0700 Subject: [PATCH 07/19] unify PlacementAsset/EmbodimentBase Signed-off-by: zhx06 --- isaaclab_arena/embodiments/droid/droid.py | 15 ++- isaaclab_arena/embodiments/embodiment_base.py | 72 ++++++++++-- .../environments/relation_solver_interface.py | 53 ++------- isaaclab_arena/relations/object_placer.py | 4 - isaaclab_arena/relations/placement_asset.py | 16 ++- isaaclab_arena/relations/placement_events.py | 42 ++----- isaaclab_arena/terms/events.py | 50 ++++++++ isaaclab_arena/tests/dummy_embodiment.py | 45 +++++++- isaaclab_arena/tests/dummy_object.py | 41 ++++++- isaaclab_arena/tests/test_placement_events.py | 109 +++++++++--------- .../tests/test_relation_solver_embodiment.py | 14 ++- .../tests/test_relation_solver_interface.py | 15 ++- 12 files changed, 300 insertions(+), 176 deletions(-) diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index eaffa1c1eb..05e8b8f386 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -41,7 +41,7 @@ from isaaclab_arena.embodiments.embodiment_base import EmbodimentBase from isaaclab_arena.embodiments.franka.franka import franka_stack_events 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, translate_by_xyz_offset # The base stand's x/y footprint. _STAND_FOOTPRINT_SCALE_XY: tuple[float, float] = (1.2, 1.2) @@ -98,9 +98,16 @@ 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 _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 diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index bbf46b89fb..59ddfe2928 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -9,6 +9,7 @@ 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.embodiments.common.arm_mode import ArmMode @@ -57,6 +58,8 @@ def __init__( 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.""" @@ -94,14 +97,53 @@ def get_collision_mesh(self) -> trimesh.Trimesh | None: self._collision_mesh = extract_trimesh_from_prim(stage, default_prim.GetPath().pathString, scale) return self._collision_mesh - def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: - """Set the embodiment root pose.""" - assert isinstance(pose, Pose), "Embodiments require one root Pose" + 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_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 supports_per_env_initial_pose(self) -> bool: - """Return False because embodiment configs store one root pose.""" - return False + 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.""" @@ -110,7 +152,7 @@ def set_joint_initial_pos(self, joint_pos: Mapping[str, float]) -> None: 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: @@ -124,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( @@ -159,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 diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py index dd048a1cc0..d4b0535c93 100644 --- a/isaaclab_arena/environments/relation_solver_interface.py +++ b/isaaclab_arena/environments/relation_solver_interface.py @@ -11,11 +11,7 @@ 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_pose_from_layout, - place_assets_from_layouts, - 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 PosePerEnv @@ -154,23 +150,22 @@ def _apply_relation_placement_result( anchor_assets=anchor_assets, ) - # Objects can store PosePerEnv, so their reset events restore fixed per-env - # poses. Scenes containing any asset without per-env pose support use one - # coordinated event to restore layouts[env_id]. - if all(asset.supports_per_env_initial_pose() for asset in assets): - _apply_static_initial_poses( - assets=assets, - placement_pool=placement_pool, - anchor_assets=anchor_assets, - num_envs=num_envs, - ) - return None - return _apply_static_spawn_pose( + # 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( assets=assets, placement_pool=placement_pool, 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( @@ -196,30 +191,6 @@ def _apply_dynamic_spawn_pose( ) -def _apply_static_spawn_pose( - assets: list[PlacementAsset], - placement_pool: PooledObjectPlacer, - anchor_assets: set[PlacementAsset], - num_envs: int, -) -> EventTermCfg: - """Return a coordinated reset event that restores one fixed layout per environment.""" - from isaaclab.managers import EventTermCfg - - layouts = placement_pool.sample_with_replacement(num_envs) - # Scene configs hold one pose per asset, so construction uses one valid layout. - # The reset event installs layouts[env_id] before the first environment step. - construction_layout = layouts[0] - _seed_spawn_config_from_layout(assets, anchor_assets, construction_layout) - return EventTermCfg( - func=place_assets_from_layouts, - mode="reset", - params={ - "assets": assets, - "layouts": layouts, - }, - ) - - def _seed_spawn_config_from_layout( assets: list[PlacementAsset], anchor_assets: set[PlacementAsset], diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py index d3a194df90..b6406a1ac3 100644 --- a/isaaclab_arena/relations/object_placer.py +++ b/isaaclab_arena/relations/object_placer.py @@ -731,10 +731,6 @@ def _yaw_delta(env_idx: int) -> float: else: obj.set_initial_pose(Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw)) else: - assert obj.supports_per_env_initial_pose(), ( - f"Placement entity '{obj.name}' cannot store per-environment poses. " - "Set apply_positions_to_objects=False and apply the results at reset." - ) poses = [ Pose( position_xyz=positions_per_env[env_idx][obj], diff --git a/isaaclab_arena/relations/placement_asset.py b/isaaclab_arena/relations/placement_asset.py index 0d1c20b619..e6c188f361 100644 --- a/isaaclab_arena/relations/placement_asset.py +++ b/isaaclab_arena/relations/placement_asset.py @@ -58,8 +58,8 @@ def get_initial_pose(self) -> Pose | PoseRange | PosePerEnv | None: def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: """Set the configured root pose. - ``PoseRange`` and ``PosePerEnv`` support is subclass-specific. Callers - assigning ``PosePerEnv`` must first check ``supports_per_env_initial_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 @@ -67,14 +67,18 @@ 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 - def supports_per_env_initial_pose(self) -> bool: - """Return whether one configured pose may be stored per environment.""" - return True - @abstractmethod def get_bounding_box(self) -> AxisAlignedBoundingBox: """Return root-relative axis-aligned bounds.""" diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py index 79e619fdb0..dc8d122d52 100644 --- a/isaaclab_arena/relations/placement_events.py +++ b/isaaclab_arena/relations/placement_events.py @@ -101,12 +101,13 @@ def write_layout_to_sim( for asset in result.positions: if asset in anchor_assets: continue - scene_asset = env.scene[asset.get_scene_name()] - pose = get_pose_from_layout(asset, result) - pose_tensor = pose.to_tensor(device=env.device).unsqueeze(0) - pose_tensor[0, :3] += env.scene.env_origins[env_id, :] - scene_asset.write_root_pose_to_sim(pose_tensor, env_ids=env_id_tensor) - scene_asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_id_tensor) + layout_pose = get_pose_from_layout(asset, result) + for scene_name, pose in asset.layout_pose_to_scene_writes(layout_pose): + scene_asset = env.scene[scene_name] + pose_tensor = pose.to_tensor(device=env.device).unsqueeze(0) + pose_tensor[0, :3] += env.scene.env_origins[env_id, :] + scene_asset.write_root_pose_to_sim(pose_tensor, env_ids=env_id_tensor) + scene_asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_id_tensor) def solve_and_place_objects( @@ -147,32 +148,3 @@ def solve_and_place_objects( ) # Only write non-anchor assets to the sim. write_layout_to_sim(env, cur_env, result, anchor_assets, base_rotations) - - -def place_assets_from_layouts( - env: ManagerBasedEnv, - env_ids: torch.Tensor | None, - assets: list[PlacementAsset], - layouts: list[PlacementResult], -) -> None: - """Restore one fixed placement layout per environment. - - Args: - env: The Isaac Lab environment. - env_ids: Environment indices to restore. - assets: Assets participating in relation solving. - layouts: Fixed layout indexed by environment. - """ - if env_ids is None or len(env_ids) == 0: - return - assert len(layouts) == env.scene.env_origins.shape[0], "Static layouts must match the scene environment count" - anchor_assets = set(get_anchor_objects(assets)) - base_rotations = get_base_rotation_per_asset(assets) - for env_id in env_ids.tolist(): - write_layout_to_sim( - env, - env_id, - layouts[env_id], - anchor_assets, - base_rotations, - ) diff --git a/isaaclab_arena/terms/events.py b/isaaclab_arena/terms/events.py index e0e859667d..48409fe491 100644 --- a/isaaclab_arena/terms/events.py +++ b/isaaclab_arena/terms/events.py @@ -63,6 +63,56 @@ 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] + asset.write_root_pose_to_sim(pose_t_xyz_q_xyzw, env_ids=env_ids) + asset.write_root_velocity_to_sim(torch.zeros(num_envs, 6, device=env.device), env_ids=env_ids) + + +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 index 55257a19fc..4f83633d6e 100644 --- a/isaaclab_arena/tests/dummy_embodiment.py +++ b/isaaclab_arena/tests/dummy_embodiment.py @@ -8,14 +8,22 @@ 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 +from isaaclab_arena.utils.pose import Pose, PosePerEnv + +if TYPE_CHECKING: + from isaaclab.managers import EventTermCfg class DummyEmbodiment(PlacementAsset): - """Embodiment geometry without simulator dependencies.""" + """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, @@ -30,6 +38,35 @@ def __init__( 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.""" @@ -39,10 +76,6 @@ def get_collision_mesh(self) -> trimesh.Trimesh | None: """Return the configured collision mesh.""" return self.collision_mesh - def supports_per_env_initial_pose(self) -> bool: - """Return False because the dummy stores one root pose.""" - return False - 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 index 664a3dcf67..0977931360 100644 --- a/isaaclab_arena/tests/dummy_object.py +++ b/isaaclab_arena/tests/dummy_object.py @@ -6,15 +6,23 @@ 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 +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.""" + """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, @@ -31,6 +39,35 @@ def __init__( 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).""" diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index 298f26a0ce..7275f5ffed 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -230,68 +230,69 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: env._assets["robot"].write_root_pose_to_sim.assert_called_once() -def test_static_layout_event_writes_embodiment_to_configured_scene_asset(): - from isaaclab_arena.relations.placement_events import place_assets_from_layouts - from isaaclab_arena.relations.placement_result import PlacementResult - from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +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 - 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", - ) - layouts = [ - PlacementResult( - validation_results=_checklist(True), - positions={robot: (0.1, 0.2, 0.0)}, - final_loss=0.0, - attempts=1, - ), - PlacementResult( - validation_results=_checklist(True), - positions={robot: (0.4, 0.5, 0.0)}, - final_loss=0.0, - attempts=1, - ), - ] 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)))] - place_assets_from_layouts( - env, - torch.tensor([1]), - assets=[robot], - layouts=layouts, - ) + reset_placement_asset_pose(env, torch.tensor([0, 1]), scene_writes=scene_writes) - assert "droid" not in env._assets - pose = env._assets["robot"].write_root_pose_to_sim.call_args.args[0] - assert torch.allclose(pose[0, :3], torch.tensor([0.4, 0.5, 0.0])) - - with pytest.raises(AssertionError, match="Static layouts must match"): - place_assets_from_layouts( - env, - torch.tensor([0]), - assets=[robot], - layouts=layouts[:1], - ) + 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_static_layout_event_rejects_missing_non_anchor_assets(): - from isaaclab_arena.relations.placement_events import place_assets_from_layouts - from isaaclab_arena.relations.placement_result import PlacementResult +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 - _, box1, _ = _create_test_objects() - env = _make_mock_env(num_envs=1) - layout = PlacementResult( - validation_results=_checklist(True), - positions={}, - final_loss=0.0, - attempts=1, - ) + 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) + 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"].write_root_pose_to_sim.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="missing non-anchor assets.*box1"): - place_assets_from_layouts(env, torch.tensor([0]), assets=[box1], layouts=[layout]) + 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(): diff --git a/isaaclab_arena/tests/test_relation_solver_embodiment.py b/isaaclab_arena/tests/test_relation_solver_embodiment.py index 9c64617161..9fc2484097 100644 --- a/isaaclab_arena/tests/test_relation_solver_embodiment.py +++ b/isaaclab_arena/tests/test_relation_solver_embodiment.py @@ -7,15 +7,13 @@ import torch -import pytest - 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 +from isaaclab_arena.utils.pose import Pose, PosePerEnv def _make_floor_and_robot(): @@ -49,11 +47,15 @@ def test_relation_solver_places_embodiment(): assert robot.get_initial_pose() is not None -def test_batched_embodiment_placement_requires_runtime_application(): +def test_batched_embodiment_placement_stores_per_env_poses(): floor, robot = _make_floor_and_robot() - with pytest.raises(AssertionError, match="cannot store per-environment poses"): - ObjectPlacer(ObjectPlacerParams(placement_seed=3)).place([floor, robot], num_envs=2) + 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(): diff --git a/isaaclab_arena/tests/test_relation_solver_interface.py b/isaaclab_arena/tests/test_relation_solver_interface.py index 50b30e5592..7949352207 100644 --- a/isaaclab_arena/tests/test_relation_solver_interface.py +++ b/isaaclab_arena/tests/test_relation_solver_interface.py @@ -154,12 +154,12 @@ def test_dynamic_spawn_pose_event_params_use_runtime_assets(): assert "placement_pool" in event_cfg.params -def test_static_embodiment_placement_uses_coordinated_reset(): +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 Pose + from isaaclab_arena.utils.pose import PosePerEnv desk = _make_desk() robot = DummyEmbodiment( @@ -181,13 +181,12 @@ def test_static_embodiment_placement_uses_coordinated_reset(): 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, Pose) - assert initial_pose.position_xyz == (0.1, 0.2, 0.0) - assert event_cfg is not None - runtime_robot = event_cfg.params["assets"][1] - assert runtime_robot in event_cfg.params["layouts"][0].positions - assert event_cfg.params["layouts"][1].positions[runtime_robot] == (0.3, 0.4, 0.0) + 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 0dca1e04d997e5a7164db044affa46db340d5ded Mon Sep 17 00:00:00 2001 From: zhx06 Date: Thu, 23 Jul 2026 22:36:00 -0700 Subject: [PATCH 08/19] fix curobo CI import errors Signed-off-by: zhx06 --- isaaclab_arena_curobo/ik_reachability_validator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 From ee9f6592775c66e5e65838f714d0126e084e4e51 Mon Sep 17 00:00:00 2001 From: zhx06 Date: Thu, 23 Jul 2026 23:09:45 -0700 Subject: [PATCH 09/19] fix dummy object impportss Signed-off-by: zhx06 --- isaaclab_arena_curobo/tests/test_ik_reachability_validator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From 6e01603e1ec1f94003eaa490c4c4d2b2fbec94bd Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Thu, 23 Jul 2026 17:14:18 +0800 Subject: [PATCH 10/19] Add kitchen task yaml with robot placement --- ...oid_pick_and_place_lightwheel_kitchen.yaml | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 isaaclab_arena_environments/droid_pick_and_place_lightwheel_kitchen.yaml 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 From 4b7f4841f0a3d385bf8fd4c6c28215bbc1a28a37 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Thu, 23 Jul 2026 17:44:51 +0800 Subject: [PATCH 11/19] Droid returns stand bbox --- isaaclab_arena/embodiments/droid/droid.py | 19 ++++++ .../tests/test_embodiment_placement_bbox.py | 65 +++++++++++++++++++ 2 files changed, 84 insertions(+) create mode 100644 isaaclab_arena/tests/test_embodiment_placement_bbox.py diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index 05e8b8f386..d26500af96 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -40,6 +40,7 @@ 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, PosePerEnv, translate_by_xyz_offset @@ -109,6 +110,24 @@ 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 root-relative bounds from the stand geometry (robot root is the placement origin). + + The relation solver supplies unlifted poses; ``set_initial_pose`` then applies + ``_robot_base_z_offset``. Shift the stand footprint by the configured stand + offset plus that z lift so ``On`` placement stays aligned with the spawned base. + """ + 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" + stand = self.scene_config.stand + spawn = 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) + + return stand_bbox.translated((0.0, 0.0, self._robot_base_z_offset)) + 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. diff --git a/isaaclab_arena/tests/test_embodiment_placement_bbox.py b/isaaclab_arena/tests/test_embodiment_placement_bbox.py new file mode 100644 index 0000000000..ea64cc5d4a --- /dev/null +++ b/isaaclab_arena/tests/test_embodiment_placement_bbox.py @@ -0,0 +1,65 @@ +# 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 torch +import traceback + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + +_KITCHEN_STAND_HEIGHT_M = 0.8 + + +def _test_droid_placement_bbox_uses_stand(simulation_app) -> bool: + """Check Droid placement bounds come from the stand footprint, not the robot mesh.""" + + from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment + from isaaclab_arena.utils.pose import Pose + 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 = (0.0, 0.0, embodiment._robot_base_z_offset) + expected = stand_bbox.translated(stand_offset) + + assert bbox.min_point.shape == expected.min_point.shape + assert torch.allclose(bbox.min_point, expected.min_point) + assert torch.allclose(bbox.max_point, expected.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.min_point[0, 2].item() - stand_offset[2]) + assert abs(natural_stand_bottom_z - floor_top_z) < 1e-3 + + except Exception as exc: + print(f"Error: {exc}") + traceback.print_exc() + return False + + return True + + +def test_droid_placement_bbox_uses_stand(): + """Pytest entry point for the Droid stand placement-bbox test.""" + result = run_simulation_app_function(_test_droid_placement_bbox_uses_stand, headless=True) + assert result, f"Test {test_droid_placement_bbox_uses_stand.__name__} failed" + + +if __name__ == "__main__": + test_droid_placement_bbox_uses_stand() From 9d391a8e5174712b34fd8a961e6af6d8c8b29abb Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Thu, 23 Jul 2026 18:52:20 +0800 Subject: [PATCH 12/19] Fix Droid robot and stand placement on reset Relation placement reset wrote only the robot root at the solver's unlifted z, leaving the stand at spawn and causing float when stand_height_m differs from the default. Map layout poses to all scene roots via layout_pose_to_scene_writes and apply stand-height lift for Droid on reset. Signed-off-by: Qian Lin --- isaaclab_arena/embodiments/droid/droid.py | 35 +++++-- isaaclab_arena/relations/placement_events.py | 46 +++++++-- .../tests/test_droid_reset_placement.py | 57 +++++++++++ .../tests/test_droid_stand_height.py | 12 +++ .../tests/test_embodiment_placement_bbox.py | 2 +- isaaclab_arena/tests/test_placement_events.py | 98 +++++++++++++++++++ 6 files changed, 235 insertions(+), 15 deletions(-) create mode 100644 isaaclab_arena/tests/test_droid_reset_placement.py diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index d26500af96..2acb218d63 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -42,10 +42,13 @@ 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, PosePerEnv, 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) +_STAND_SCENE_NAME: str = "stand" # The default stand height. _DEFAULT_STAND_HEIGHT_M: float = 1.35 _FALLBACK_STAND_UNIT_HEIGHT_M: float = 0.795 @@ -114,8 +117,8 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox: """Return root-relative bounds from the stand geometry (robot root is the placement origin). The relation solver supplies unlifted poses; ``set_initial_pose`` then applies - ``_robot_base_z_offset``. Shift the stand footprint by the configured stand - offset plus that z lift so ``On`` placement stays aligned with the spawned base. + ``_robot_base_offset``. Shift the stand footprint by that z lift so ``On`` + placement stays aligned with the spawned base. """ from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd @@ -126,14 +129,30 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox: scale = tuple(spawn.scale or (1.0, 1.0, 1.0)) stand_bbox = compute_local_bounding_box_from_usd(spawn.usd_path, scale) - return stand_bbox.translated((0.0, 0.0, self._robot_base_z_offset)) + return stand_bbox.translated((0.0, 0.0, self._robot_base_offset[2])) + + 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_SCENE_NAME, 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: diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py index dc8d122d52..6ed2f12012 100644 --- a/isaaclab_arena/relations/placement_events.py +++ b/isaaclab_arena/relations/placement_events.py @@ -73,6 +73,37 @@ def get_movable_asset_names( 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, + zero_velocity: 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_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_id_tensor) + scene_asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_id_tensor) + 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_id_tensor.detach().cpu(), + ) + return + + assert False, f"Scene asset '{scene_name}' does not support root pose writes" + + def write_layout_to_sim( env: ManagerBasedEnv, env_id: int, @@ -102,12 +133,15 @@ def write_layout_to_sim( if asset in anchor_assets: continue layout_pose = get_pose_from_layout(asset, result) - for scene_name, pose in asset.layout_pose_to_scene_writes(layout_pose): - scene_asset = env.scene[scene_name] - pose_tensor = pose.to_tensor(device=env.device).unsqueeze(0) - pose_tensor[0, :3] += env.scene.env_origins[env_id, :] - scene_asset.write_root_pose_to_sim(pose_tensor, env_ids=env_id_tensor) - scene_asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_id_tensor) + 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, + zero_velocity, + ) def solve_and_place_objects( diff --git a/isaaclab_arena/tests/test_droid_reset_placement.py b/isaaclab_arena/tests/test_droid_reset_placement.py new file mode 100644 index 0000000000..d31c4ce475 --- /dev/null +++ b/isaaclab_arena/tests/test_droid_reset_placement.py @@ -0,0 +1,57 @@ +# 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 + +"""Verify Droid robot and stand stay aligned after relation placement reset.""" + +from __future__ import annotations + +import traceback +from pathlib import Path + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + +_KITCHEN_YAML = Path(__file__).resolve().parents[2] / "isaaclab_arena_environments" / "kitchen_task.yaml" +_STAND_HEIGHT_M = 0.8 +_Z_MATCH_EPS = 1e-3 + + +def _test_droid_reset_placement(simulation_app) -> bool: + 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: + import warp as wp + + spec = ArenaEnvGraphSpec.from_yaml(_KITCHEN_YAML) + params = dict(spec.embodiment.params) + params["stand_height_m"] = _STAND_HEIGHT_M + spec.embodiment.params = params + + builder = ArenaEnvBuilder(spec.to_arena_env(), ArenaEnvBuilderCfg(num_envs=1)) + env = builder.make_registered() + env.reset() + + robot_z = wp.to_torch(env.unwrapped.scene["robot"].data.root_link_pose_w)[0, 2].item() + stand_pos = env.unwrapped.scene["stand"].get_world_poses()[0].torch + stand_z = stand_pos[0, 2].item() + assert abs(robot_z - stand_z) < _Z_MATCH_EPS, f"robot z {robot_z} != stand z {stand_z} after reset" + + env.close() + except Exception as exc: + print(f"Error: {exc}") + traceback.print_exc() + return False + + return True + + +def test_droid_reset_placement(): + result = run_simulation_app_function(_test_droid_reset_placement, headless=True) + assert result, f"Test {test_droid_reset_placement.__name__} failed" + + +if __name__ == "__main__": + test_droid_reset_placement() diff --git a/isaaclab_arena/tests/test_droid_stand_height.py b/isaaclab_arena/tests/test_droid_stand_height.py index 300d40c2d9..05e09e2ae6 100644 --- a/isaaclab_arena/tests/test_droid_stand_height.py +++ b/isaaclab_arena/tests/test_droid_stand_height.py @@ -55,10 +55,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) diff --git a/isaaclab_arena/tests/test_embodiment_placement_bbox.py b/isaaclab_arena/tests/test_embodiment_placement_bbox.py index ea64cc5d4a..277db0fa24 100644 --- a/isaaclab_arena/tests/test_embodiment_placement_bbox.py +++ b/isaaclab_arena/tests/test_embodiment_placement_bbox.py @@ -24,7 +24,7 @@ def _test_droid_placement_bbox_uses_stand(simulation_app) -> bool: stand = embodiment.scene_config.stand stand_bbox = compute_local_bounding_box_from_usd(stand.spawn.usd_path, tuple(stand.spawn.scale)) - stand_offset = (0.0, 0.0, embodiment._robot_base_z_offset) + stand_offset = embodiment._robot_base_offset expected = stand_bbox.translated(stand_offset) assert bbox.min_point.shape == expected.min_point.shape diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index 7275f5ffed..b17d11aff3 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -255,6 +255,104 @@ def test_reset_placement_asset_pose_writes_compound_prims_with_env_origins(): assert torch.count_nonzero(velocity) == 0 +def test_write_layout_to_sim_writes_companion_scene_assets(): + from isaaclab_arena.relations.placement_events import write_layout_to_sim + from isaaclab_arena.relations.placement_result import PlacementResult + from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.pose import Pose + + class CompanionEmbodiment(DummyEmbodiment): + def layout_pose_to_scene_writes(self, layout_pose: Pose) -> list[tuple[str, Pose]]: + companion_pose = Pose( + position_xyz=( + layout_pose.position_xyz[0] + 1.0, + layout_pose.position_xyz[1], + layout_pose.position_xyz[2], + ), + rotation_xyzw=layout_pose.rotation_xyzw, + ) + return [ + (self.get_scene_name(), layout_pose), + ("companion", companion_pose), + ] + + robot = CompanionEmbodiment( + 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) + layout = PlacementResult( + validation_results=_checklist(True), + positions={robot: (0.2, 0.3, 0.4)}, + final_loss=0.0, + attempts=1, + ) + + write_layout_to_sim(env, 0, layout, anchor_assets=set(), base_rotations={robot: (0.0, 0.0, 0.0, 1.0)}) + + robot_pose = env._assets["robot"].write_root_pose_to_sim.call_args.args[0] + companion_pose = env._assets["companion"].write_root_pose_to_sim.call_args.args[0] + assert torch.allclose(robot_pose[0, :3], torch.tensor([0.2, 0.3, 0.4])) + assert torch.allclose(companion_pose[0, :3], torch.tensor([1.2, 0.3, 0.4])) + + +def test_solve_and_place_objects_writes_droid_robot_and_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) + + 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_pose = env._assets["stand"].write_root_pose_to_sim.call_args.args[0] + expected_robot_z = 1.36 + expected_offset + assert abs(robot_pose[0, 2].item() - expected_robot_z) < 1e-5 + assert abs(stand_pose[0, 2].item() - expected_robot_z) < 1e-5 + assert abs(stand_pose[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 From d69ded3ca09445c2e00ad2491b38ad165c54744d Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Fri, 24 Jul 2026 15:38:35 +0800 Subject: [PATCH 13/19] Address self review comments * Clean up code comments * Fix bug in droid get_bounding_box offset handling --- isaaclab_arena/embodiments/droid/droid.py | 19 ++++++++----------- isaaclab_arena/relations/placement_events.py | 2 ++ 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index 2acb218d63..dae1c9a24d 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -48,7 +48,6 @@ _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) -_STAND_SCENE_NAME: str = "stand" # The default stand height. _DEFAULT_STAND_HEIGHT_M: float = 1.35 _FALLBACK_STAND_UNIT_HEIGHT_M: float = 0.795 @@ -114,22 +113,20 @@ def set_spawn_pose(self, pose: Pose) -> None: super().set_spawn_pose(pose.translate(self._robot_base_offset)) def get_bounding_box(self) -> AxisAlignedBoundingBox: - """Return root-relative bounds from the stand geometry (robot root is the placement origin). + """Return stand bounding box as proxy for the robot bounding box. - The relation solver supplies unlifted poses; ``set_initial_pose`` then applies - ``_robot_base_offset``. Shift the stand footprint by that z lift so ``On`` - placement stays aligned with the spawned base. + 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" - stand = self.scene_config.stand - spawn = stand.spawn + 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) - - return stand_bbox.translated((0.0, 0.0, self._robot_base_offset[2])) + 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.""" @@ -144,7 +141,7 @@ def layout_pose_to_scene_writes(self, layout_pose: Pose) -> list[tuple[str, Pose stand_pose = self._stand_pose_from_robot_pose(robot_pose) return [ (self.get_scene_name(), robot_pose), - (_STAND_SCENE_NAME, stand_pose), + ("stand", stand_pose), ] def _update_scene_cfg_with_robot_initial_pose(self, scene_config: Any, pose: Pose) -> Any: @@ -312,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/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py index 6ed2f12012..2e2b5bb62f 100644 --- a/isaaclab_arena/relations/placement_events.py +++ b/isaaclab_arena/relations/placement_events.py @@ -86,12 +86,14 @@ def _write_scene_root_pose_to_sim( pose_tensor = sim_pose.to_tensor(device=env.device).unsqueeze(0) pose_tensor[0, :3] += env.scene.env_origins[env_id, :] + # Articulations and rigid objects use write_root_pose_to_sim (ArticulationCfg/RigidObjectSetCfg) 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_id_tensor) scene_asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_id_tensor) return + # AssetBase extras (e.g. Droid stand) use set_world_poses (AssetBaseCfg). set_world_poses = getattr(scene_asset, "set_world_poses", None) if set_world_poses is not None: set_world_poses( From 882a0cfb28243a6667d147a6a802dca9d2bdf501 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Fri, 24 Jul 2026 15:49:53 +0800 Subject: [PATCH 14/19] Consolidate Droid stand and placement tests into one module. Merge placement-bbox and reset-alignment coverage into test_droid_stand_height.py so the suite uses one sim startup and drops redundant cases. Remove the overlapping companion-scene write test; the Droid reset write test covers multi-prim layout_pose_to_scene_writes. Signed-off-by: Qian Lin --- .../tests/test_droid_reset_placement.py | 57 ----- .../tests/test_droid_stand_height.py | 202 ++++++++++++------ .../tests/test_embodiment_placement_bbox.py | 65 ------ isaaclab_arena/tests/test_placement_events.py | 44 +--- 4 files changed, 140 insertions(+), 228 deletions(-) delete mode 100644 isaaclab_arena/tests/test_droid_reset_placement.py delete mode 100644 isaaclab_arena/tests/test_embodiment_placement_bbox.py diff --git a/isaaclab_arena/tests/test_droid_reset_placement.py b/isaaclab_arena/tests/test_droid_reset_placement.py deleted file mode 100644 index d31c4ce475..0000000000 --- a/isaaclab_arena/tests/test_droid_reset_placement.py +++ /dev/null @@ -1,57 +0,0 @@ -# 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 - -"""Verify Droid robot and stand stay aligned after relation placement reset.""" - -from __future__ import annotations - -import traceback -from pathlib import Path - -from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function - -_KITCHEN_YAML = Path(__file__).resolve().parents[2] / "isaaclab_arena_environments" / "kitchen_task.yaml" -_STAND_HEIGHT_M = 0.8 -_Z_MATCH_EPS = 1e-3 - - -def _test_droid_reset_placement(simulation_app) -> bool: - 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: - import warp as wp - - spec = ArenaEnvGraphSpec.from_yaml(_KITCHEN_YAML) - params = dict(spec.embodiment.params) - params["stand_height_m"] = _STAND_HEIGHT_M - spec.embodiment.params = params - - builder = ArenaEnvBuilder(spec.to_arena_env(), ArenaEnvBuilderCfg(num_envs=1)) - env = builder.make_registered() - env.reset() - - robot_z = wp.to_torch(env.unwrapped.scene["robot"].data.root_link_pose_w)[0, 2].item() - stand_pos = env.unwrapped.scene["stand"].get_world_poses()[0].torch - stand_z = stand_pos[0, 2].item() - assert abs(robot_z - stand_z) < _Z_MATCH_EPS, f"robot z {robot_z} != stand z {stand_z} after reset" - - env.close() - except Exception as exc: - print(f"Error: {exc}") - traceback.print_exc() - return False - - return True - - -def test_droid_reset_placement(): - result = run_simulation_app_function(_test_droid_reset_placement, headless=True) - assert result, f"Test {test_droid_reset_placement.__name__} failed" - - -if __name__ == "__main__": - test_droid_reset_placement() diff --git a/isaaclab_arena/tests/test_droid_stand_height.py b/isaaclab_arena/tests/test_droid_stand_height.py index 05e09e2ae6..b5d39c24e5 100644 --- a/isaaclab_arena/tests/test_droid_stand_height.py +++ b/isaaclab_arena/tests/test_droid_stand_height.py @@ -3,90 +3,166 @@ # # SPDX-License-Identifier: Apache-2.0 +"""Droid stand height, placement bbox, and reset alignment tests.""" + +from __future__ import annotations + +import torch import traceback +from collections.abc import Callable +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 -def _test_droid_stand_height(simulation_app) -> bool: - """Check ``stand_height_m`` (absolute meters) sets the stand z-scale and lifts the robot base.""" - - from isaaclab_arena.assets.registries import AssetRegistry +def _check_stand_height_scaling(simulation_app) -> None: + """``stand_height_m`` sets stand z-scale and lifts the default robot/stand base together.""" from isaaclab_arena.embodiments.droid.droid import ( _DEFAULT_STAND_HEIGHT_M, _STAND_FOOTPRINT_SCALE_XY, DroidAbsoluteJointPositionEmbodiment, _stand_unit_height_m, ) + + default_emb = DroidAbsoluteJointPositionEmbodiment() + unit_height = _stand_unit_height_m(default_emb.scene_config.stand.spawn.usd_path) + expected_default_scale = (*_STAND_FOOTPRINT_SCALE_XY, _DEFAULT_STAND_HEIGHT_M / unit_height) + expected_custom_scale = (*_STAND_FOOTPRINT_SCALE_XY, _CUSTOM_STAND_HEIGHT_M / unit_height) + expected_offset = _CUSTOM_STAND_HEIGHT_M - _DEFAULT_STAND_HEIGHT_M + + for got, want in zip(default_emb.scene_config.stand.spawn.scale, expected_default_scale): + assert abs(got - want) < 1e-6 + assert default_emb.scene_config.robot.init_state.pos[2] == 0.0 + + custom_emb = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_CUSTOM_STAND_HEIGHT_M) + for got, want in zip(custom_emb.scene_config.stand.spawn.scale, expected_custom_scale): + assert abs(got - want) < 1e-6 + assert custom_emb.scene_config.robot.spawn.scale in (None, (1.0, 1.0, 1.0)) + assert abs(custom_emb.scene_config.robot.init_state.pos[2] - expected_offset) < 1e-6 + assert abs(custom_emb.scene_config.stand.init_state.pos[2] - expected_offset) < 1e-6 + + +def _check_robot_stand_pose_wiring(simulation_app) -> None: + """Lifted robot poses drive stand scene cfg and reset scene writes via the stand root offset.""" + from isaaclab_arena.embodiments.droid.droid import ( + _DEFAULT_STAND_HEIGHT_M, + _STAND_ROOT_OFFSET_IN_ROBOT_FRAME, + DroidAbsoluteJointPositionEmbodiment, + ) from isaaclab_arena.utils.pose import Pose + expected_offset = _CUSTOM_STAND_HEIGHT_M - _DEFAULT_STAND_HEIGHT_M + posed_emb = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_CUSTOM_STAND_HEIGHT_M) + posed_emb.set_initial_pose(Pose(position_xyz=(0.3, 0.0, 0.5), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + assert abs(posed_emb.initial_pose.position_xyz[2] - (0.5 + expected_offset)) < 1e-6 + + scene_cfg = posed_emb.get_scene_cfg() + assert tuple(posed_emb.initial_pose.position_xyz) == tuple(scene_cfg.robot.init_state.pos) + assert abs(scene_cfg.stand.init_state.pos[0] - (0.3 + _STAND_ROOT_OFFSET_IN_ROBOT_FRAME[0])) < 1e-6 + assert abs(scene_cfg.stand.init_state.pos[2] - scene_cfg.robot.init_state.pos[2]) < 1e-6 + + 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 scene_writes[0][0] == "robot" + assert scene_writes[1][0] == "stand" + robot_write_pose = scene_writes[0][1] + stand_write_pose = scene_writes[1][1] + assert ( + abs( + stand_write_pose.position_xyz[0] - (robot_write_pose.position_xyz[0] + _STAND_ROOT_OFFSET_IN_ROBOT_FRAME[0]) + ) + < 1e-6 + ) + assert abs(stand_write_pose.position_xyz[2] - robot_write_pose.position_xyz[2]) < 1e-6 + + +def _check_placement_bbox(simulation_app) -> None: + """Placement bounds follow the stand footprint, including stand root and stand-height offsets.""" + 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 + + kitchen_emb = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_KITCHEN_STAND_HEIGHT_M) + bbox = kitchen_emb.get_bounding_box() + stand = kitchen_emb.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, kitchen_emb._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( + kitchen_emb.scene_config.robot.spawn.usd_path, + tuple(kitchen_emb.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 + + floor_top_z = 0.0 + solver_z = floor_top_z - bbox.min_point[0, 2].item() + kitchen_emb.set_initial_pose(Pose(position_xyz=(0.0, 0.0, solver_z), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + spawn_z = kitchen_emb.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 + + +def _check_kitchen_reset_alignment(simulation_app) -> None: + """Robot and stand roots stay z-aligned after relation placement reset in the kitchen task.""" + 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 + + 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() try: - # ``stand_height_m`` is an absolute height in meters, converted to a z-scale via the stand's - # native (scale=1.0) height. - default_emb = DroidAbsoluteJointPositionEmbodiment() - unit_height = _stand_unit_height_m(default_emb.scene_config.stand.spawn.usd_path) - expected_default_scale = (*_STAND_FOOTPRINT_SCALE_XY, _DEFAULT_STAND_HEIGHT_M / unit_height) - expected_custom_scale = (*_STAND_FOOTPRINT_SCALE_XY, _CUSTOM_STAND_HEIGHT_M / unit_height) - - # The default leaves the robot base at z=0 (no lift relative to the default height). - for got, want in zip(default_emb.scene_config.stand.spawn.scale, expected_default_scale): - assert abs(got - want) < 1e-6 - assert default_emb.scene_config.robot.init_state.pos[2] == 0.0 - - # The lift is the height delta from the default, in meters. - expected_offset = _CUSTOM_STAND_HEIGHT_M - _DEFAULT_STAND_HEIGHT_M - - # An override changes only the z-scale (x/y footprint and robot mesh untouched)... - custom_emb = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_CUSTOM_STAND_HEIGHT_M) - for got, want in zip(custom_emb.scene_config.stand.spawn.scale, expected_custom_scale): - assert abs(got - want) < 1e-6 - assert custom_emb.scene_config.robot.spawn.scale in (None, (1.0, 1.0, 1.0)) - - # ...and lifts the robot base and stand together so the stand's floor contact is preserved. - assert abs(custom_emb.scene_config.robot.init_state.pos[2] - expected_offset) < 1e-6 - assert abs(custom_emb.scene_config.stand.init_state.pos[2] - expected_offset) < 1e-6 - - # An explicit initial_pose is lifted at ingestion (set_initial_pose), so the stored override - # already equals the spawned base: the requested z plus the stand-height offset. - posed_emb = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_CUSTOM_STAND_HEIGHT_M) - posed_emb.set_initial_pose(Pose(position_xyz=(0.3, 0.0, 0.5), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) - 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) - for got, want in zip(registry_emb.scene_config.stand.spawn.scale, expected_custom_scale): - assert abs(got - want) < 1e-6 - - except Exception as e: - print(f"Error: {e}") - traceback.print_exc() - return False + env.reset() + robot_z = wp.to_torch(env.unwrapped.scene["robot"].data.root_link_pose_w)[0, 2].item() + stand_z = env.unwrapped.scene["stand"].get_world_poses()[0].torch[0, 2].item() + assert abs(robot_z - stand_z) < _Z_MATCH_EPS, f"robot z {robot_z} != stand z {stand_z} after reset" + finally: + env.close() + +_DROID_STAND_CHECKS: tuple[Callable[..., None], ...] = ( + _check_stand_height_scaling, + _check_robot_stand_pose_wiring, + _check_placement_bbox, + _check_kitchen_reset_alignment, +) + + +def _test_droid_stand_height(simulation_app) -> bool: + """Run Droid stand/placement checks in one SimulationApp session.""" + for check in _DROID_STAND_CHECKS: + try: + check(simulation_app) + except Exception as exc: + print(f"Error in {check.__name__}: {exc}") + traceback.print_exc() + return False return True def test_droid_stand_height(): - """Pytest entry point for the Droid stand-height configuration test.""" + """Pytest entry point for Droid stand height, placement bbox, and reset alignment.""" result = run_simulation_app_function(_test_droid_stand_height, headless=True) assert result, f"Test {test_droid_stand_height.__name__} failed" diff --git a/isaaclab_arena/tests/test_embodiment_placement_bbox.py b/isaaclab_arena/tests/test_embodiment_placement_bbox.py deleted file mode 100644 index 277db0fa24..0000000000 --- a/isaaclab_arena/tests/test_embodiment_placement_bbox.py +++ /dev/null @@ -1,65 +0,0 @@ -# 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 torch -import traceback - -from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function - -_KITCHEN_STAND_HEIGHT_M = 0.8 - - -def _test_droid_placement_bbox_uses_stand(simulation_app) -> bool: - """Check Droid placement bounds come from the stand footprint, not the robot mesh.""" - - from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment - from isaaclab_arena.utils.pose import Pose - 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 = embodiment._robot_base_offset - expected = stand_bbox.translated(stand_offset) - - assert bbox.min_point.shape == expected.min_point.shape - assert torch.allclose(bbox.min_point, expected.min_point) - assert torch.allclose(bbox.max_point, expected.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.min_point[0, 2].item() - stand_offset[2]) - assert abs(natural_stand_bottom_z - floor_top_z) < 1e-3 - - except Exception as exc: - print(f"Error: {exc}") - traceback.print_exc() - return False - - return True - - -def test_droid_placement_bbox_uses_stand(): - """Pytest entry point for the Droid stand placement-bbox test.""" - result = run_simulation_app_function(_test_droid_placement_bbox_uses_stand, headless=True) - assert result, f"Test {test_droid_placement_bbox_uses_stand.__name__} failed" - - -if __name__ == "__main__": - test_droid_placement_bbox_uses_stand() diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index b17d11aff3..2e769bcfb4 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -255,50 +255,8 @@ def test_reset_placement_asset_pose_writes_compound_prims_with_env_origins(): assert torch.count_nonzero(velocity) == 0 -def test_write_layout_to_sim_writes_companion_scene_assets(): - from isaaclab_arena.relations.placement_events import write_layout_to_sim - from isaaclab_arena.relations.placement_result import PlacementResult - from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment - from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox - from isaaclab_arena.utils.pose import Pose - - class CompanionEmbodiment(DummyEmbodiment): - def layout_pose_to_scene_writes(self, layout_pose: Pose) -> list[tuple[str, Pose]]: - companion_pose = Pose( - position_xyz=( - layout_pose.position_xyz[0] + 1.0, - layout_pose.position_xyz[1], - layout_pose.position_xyz[2], - ), - rotation_xyzw=layout_pose.rotation_xyzw, - ) - return [ - (self.get_scene_name(), layout_pose), - ("companion", companion_pose), - ] - - robot = CompanionEmbodiment( - 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) - layout = PlacementResult( - validation_results=_checklist(True), - positions={robot: (0.2, 0.3, 0.4)}, - final_loss=0.0, - attempts=1, - ) - - write_layout_to_sim(env, 0, layout, anchor_assets=set(), base_rotations={robot: (0.0, 0.0, 0.0, 1.0)}) - - robot_pose = env._assets["robot"].write_root_pose_to_sim.call_args.args[0] - companion_pose = env._assets["companion"].write_root_pose_to_sim.call_args.args[0] - assert torch.allclose(robot_pose[0, :3], torch.tensor([0.2, 0.3, 0.4])) - assert torch.allclose(companion_pose[0, :3], torch.tensor([1.2, 0.3, 0.4])) - - 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 From 7ee193e213b75a8f0a4b228e84c178e92071c1ff Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Fri, 24 Jul 2026 16:18:21 +0800 Subject: [PATCH 15/19] Split Droid placement sim tests while preserving stand-height case. Keep the original _test_droid_stand_height body and move placement bbox and kitchen reset checks into separate pytest entry points. Signed-off-by: Qian Lin --- .../tests/test_droid_stand_height.py | 237 ++++++++++-------- 1 file changed, 129 insertions(+), 108 deletions(-) diff --git a/isaaclab_arena/tests/test_droid_stand_height.py b/isaaclab_arena/tests/test_droid_stand_height.py index b5d39c24e5..e1cbeb5e20 100644 --- a/isaaclab_arena/tests/test_droid_stand_height.py +++ b/isaaclab_arena/tests/test_droid_stand_height.py @@ -3,13 +3,8 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Droid stand height, placement bbox, and reset alignment tests.""" - -from __future__ import annotations - import torch import traceback -from collections.abc import Callable from pathlib import Path from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -22,69 +17,84 @@ _Z_MATCH_EPS = 1e-3 -def _check_stand_height_scaling(simulation_app) -> None: - """``stand_height_m`` sets stand z-scale and lifts the default robot/stand base together.""" +def _test_droid_stand_height(simulation_app) -> bool: + """Check ``stand_height_m`` (absolute meters) sets the stand z-scale and lifts the robot base.""" + + from isaaclab_arena.assets.registries import AssetRegistry from isaaclab_arena.embodiments.droid.droid import ( _DEFAULT_STAND_HEIGHT_M, _STAND_FOOTPRINT_SCALE_XY, DroidAbsoluteJointPositionEmbodiment, _stand_unit_height_m, ) + from isaaclab_arena.utils.pose import Pose - default_emb = DroidAbsoluteJointPositionEmbodiment() - unit_height = _stand_unit_height_m(default_emb.scene_config.stand.spawn.usd_path) - expected_default_scale = (*_STAND_FOOTPRINT_SCALE_XY, _DEFAULT_STAND_HEIGHT_M / unit_height) - expected_custom_scale = (*_STAND_FOOTPRINT_SCALE_XY, _CUSTOM_STAND_HEIGHT_M / unit_height) - expected_offset = _CUSTOM_STAND_HEIGHT_M - _DEFAULT_STAND_HEIGHT_M - - for got, want in zip(default_emb.scene_config.stand.spawn.scale, expected_default_scale): - assert abs(got - want) < 1e-6 - assert default_emb.scene_config.robot.init_state.pos[2] == 0.0 - - custom_emb = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_CUSTOM_STAND_HEIGHT_M) - for got, want in zip(custom_emb.scene_config.stand.spawn.scale, expected_custom_scale): - assert abs(got - want) < 1e-6 - assert custom_emb.scene_config.robot.spawn.scale in (None, (1.0, 1.0, 1.0)) - assert abs(custom_emb.scene_config.robot.init_state.pos[2] - expected_offset) < 1e-6 - assert abs(custom_emb.scene_config.stand.init_state.pos[2] - expected_offset) < 1e-6 - + try: + # ``stand_height_m`` is an absolute height in meters, converted to a z-scale via the stand's + # native (scale=1.0) height. + default_emb = DroidAbsoluteJointPositionEmbodiment() + unit_height = _stand_unit_height_m(default_emb.scene_config.stand.spawn.usd_path) + expected_default_scale = (*_STAND_FOOTPRINT_SCALE_XY, _DEFAULT_STAND_HEIGHT_M / unit_height) + expected_custom_scale = (*_STAND_FOOTPRINT_SCALE_XY, _CUSTOM_STAND_HEIGHT_M / unit_height) + + # The default leaves the robot base at z=0 (no lift relative to the default height). + for got, want in zip(default_emb.scene_config.stand.spawn.scale, expected_default_scale): + assert abs(got - want) < 1e-6 + assert default_emb.scene_config.robot.init_state.pos[2] == 0.0 + + # The lift is the height delta from the default, in meters. + expected_offset = _CUSTOM_STAND_HEIGHT_M - _DEFAULT_STAND_HEIGHT_M + + # An override changes only the z-scale (x/y footprint and robot mesh untouched)... + custom_emb = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_CUSTOM_STAND_HEIGHT_M) + for got, want in zip(custom_emb.scene_config.stand.spawn.scale, expected_custom_scale): + assert abs(got - want) < 1e-6 + assert custom_emb.scene_config.robot.spawn.scale in (None, (1.0, 1.0, 1.0)) + + # ...and lifts the robot base and stand together so the stand's floor contact is preserved. + assert abs(custom_emb.scene_config.robot.init_state.pos[2] - expected_offset) < 1e-6 + assert abs(custom_emb.scene_config.stand.init_state.pos[2] - expected_offset) < 1e-6 + + # An explicit initial_pose is lifted at ingestion (set_initial_pose), so the stored override + # already equals the spawned base: the requested z plus the stand-height offset. + posed_emb = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_CUSTOM_STAND_HEIGHT_M) + posed_emb.set_initial_pose(Pose(position_xyz=(0.3, 0.0, 0.5), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + 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) + for got, want in zip(registry_emb.scene_config.stand.spawn.scale, expected_custom_scale): + assert abs(got - want) < 1e-6 + + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False -def _check_robot_stand_pose_wiring(simulation_app) -> None: - """Lifted robot poses drive stand scene cfg and reset scene writes via the stand root offset.""" - from isaaclab_arena.embodiments.droid.droid import ( - _DEFAULT_STAND_HEIGHT_M, - _STAND_ROOT_OFFSET_IN_ROBOT_FRAME, - DroidAbsoluteJointPositionEmbodiment, - ) - from isaaclab_arena.utils.pose import Pose + return True - expected_offset = _CUSTOM_STAND_HEIGHT_M - _DEFAULT_STAND_HEIGHT_M - posed_emb = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_CUSTOM_STAND_HEIGHT_M) - posed_emb.set_initial_pose(Pose(position_xyz=(0.3, 0.0, 0.5), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) - assert abs(posed_emb.initial_pose.position_xyz[2] - (0.5 + expected_offset)) < 1e-6 - - scene_cfg = posed_emb.get_scene_cfg() - assert tuple(posed_emb.initial_pose.position_xyz) == tuple(scene_cfg.robot.init_state.pos) - assert abs(scene_cfg.stand.init_state.pos[0] - (0.3 + _STAND_ROOT_OFFSET_IN_ROBOT_FRAME[0])) < 1e-6 - assert abs(scene_cfg.stand.init_state.pos[2] - scene_cfg.robot.init_state.pos[2]) < 1e-6 - - 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 scene_writes[0][0] == "robot" - assert scene_writes[1][0] == "stand" - robot_write_pose = scene_writes[0][1] - stand_write_pose = scene_writes[1][1] - assert ( - abs( - stand_write_pose.position_xyz[0] - (robot_write_pose.position_xyz[0] + _STAND_ROOT_OFFSET_IN_ROBOT_FRAME[0]) - ) - < 1e-6 - ) - assert abs(stand_write_pose.position_xyz[2] - robot_write_pose.position_xyz[2]) < 1e-6 +def _test_droid_placement_bbox(simulation_app) -> bool: + """Check Droid placement bounds come from the stand footprint, not the robot mesh.""" -def _check_placement_bbox(simulation_app) -> None: - """Placement bounds follow the stand footprint, including stand root and stand-height offsets.""" from isaaclab_arena.embodiments.droid.droid import ( _STAND_ROOT_OFFSET_IN_ROBOT_FRAME, DroidAbsoluteJointPositionEmbodiment, @@ -92,80 +102,91 @@ def _check_placement_bbox(simulation_app) -> None: from isaaclab_arena.utils.pose import Pose, translate_by_xyz_offset from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd - kitchen_emb = DroidAbsoluteJointPositionEmbodiment(stand_height_m=_KITCHEN_STAND_HEIGHT_M) - bbox = kitchen_emb.get_bounding_box() - stand = kitchen_emb.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, kitchen_emb._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( - kitchen_emb.scene_config.robot.spawn.usd_path, - tuple(kitchen_emb.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 + 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 - floor_top_z = 0.0 - solver_z = floor_top_z - bbox.min_point[0, 2].item() - kitchen_emb.set_initial_pose(Pose(position_xyz=(0.0, 0.0, solver_z), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) - spawn_z = kitchen_emb.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 +def _test_droid_kitchen_reset_alignment(simulation_app) -> bool: + """Verify Droid robot and stand stay z-aligned after relation placement reset.""" -def _check_kitchen_reset_alignment(simulation_app) -> None: - """Robot and stand roots stay z-aligned after relation placement reset in the kitchen task.""" 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 - 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() 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_z = wp.to_torch(env.unwrapped.scene["robot"].data.root_link_pose_w)[0, 2].item() stand_z = env.unwrapped.scene["stand"].get_world_poses()[0].torch[0, 2].item() assert abs(robot_z - stand_z) < _Z_MATCH_EPS, f"robot z {robot_z} != stand z {stand_z} after reset" - finally: - env.close() - - -_DROID_STAND_CHECKS: tuple[Callable[..., None], ...] = ( - _check_stand_height_scaling, - _check_robot_stand_pose_wiring, - _check_placement_bbox, - _check_kitchen_reset_alignment, -) + env.close() + except Exception as exc: + print(f"Error: {exc}") + traceback.print_exc() + return False -def _test_droid_stand_height(simulation_app) -> bool: - """Run Droid stand/placement checks in one SimulationApp session.""" - for check in _DROID_STAND_CHECKS: - try: - check(simulation_app) - except Exception as exc: - print(f"Error in {check.__name__}: {exc}") - traceback.print_exc() - return False return True def test_droid_stand_height(): - """Pytest entry point for Droid stand height, placement bbox, and reset alignment.""" + """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() From cd3c42cc4f9f82642d6ae9fe3019acf5a334d886 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Fri, 24 Jul 2026 16:28:59 +0800 Subject: [PATCH 16/19] Fix stand z read in kitchen reset alignment test get_world_poses()[0] is already a torch.Tensor; accessing .torch raised AttributeError and prevented the kitchen reset check from running. Signed-off-by: Qian Lin --- isaaclab_arena/tests/test_droid_stand_height.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/isaaclab_arena/tests/test_droid_stand_height.py b/isaaclab_arena/tests/test_droid_stand_height.py index e1cbeb5e20..3c8eb4d68e 100644 --- a/isaaclab_arena/tests/test_droid_stand_height.py +++ b/isaaclab_arena/tests/test_droid_stand_height.py @@ -156,7 +156,7 @@ def _test_droid_kitchen_reset_alignment(simulation_app) -> bool: env.reset() robot_z = wp.to_torch(env.unwrapped.scene["robot"].data.root_link_pose_w)[0, 2].item() - stand_z = env.unwrapped.scene["stand"].get_world_poses()[0].torch[0, 2].item() + stand_z = env.unwrapped.scene["stand"].get_world_poses()[0][0, 2].item() assert abs(robot_z - stand_z) < _Z_MATCH_EPS, f"robot z {robot_z} != stand z {stand_z} after reset" env.close() From 2211b8b6fc6d3b5d7671131402c7340b00583995 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Fri, 24 Jul 2026 16:31:05 +0800 Subject: [PATCH 17/19] Share stand-aware scene root write dispatch across reset paths Extract write_scene_root_poses_to_sim so relation placement and fixed-pose embodiment resets both route AssetBase extras (e.g. Droid stand) through set_world_poses instead of write_root_pose_to_sim. Signed-off-by: Qian Lin --- isaaclab_arena/relations/placement_events.py | 25 +--------- isaaclab_arena/terms/events.py | 4 +- isaaclab_arena/utils/scene_pose_writes.py | 48 ++++++++++++++++++++ 3 files changed, 52 insertions(+), 25 deletions(-) create mode 100644 isaaclab_arena/utils/scene_pose_writes.py diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py index 2e2b5bb62f..8b85d5c4cd 100644 --- a/isaaclab_arena/relations/placement_events.py +++ b/isaaclab_arena/relations/placement_events.py @@ -10,7 +10,7 @@ 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: @@ -79,31 +79,12 @@ def _write_scene_root_pose_to_sim( sim_pose: Pose, env_id: int, env_id_tensor: torch.Tensor, - zero_velocity: 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, :] - - # Articulations and rigid objects use write_root_pose_to_sim (ArticulationCfg/RigidObjectSetCfg) - 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_id_tensor) - scene_asset.write_root_velocity_to_sim(zero_velocity, env_ids=env_id_tensor) - return - - # AssetBase extras (e.g. Droid stand) use set_world_poses (AssetBaseCfg). - 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_id_tensor.detach().cpu(), - ) - return - - assert False, f"Scene asset '{scene_name}' does not support root pose writes" + write_scene_root_poses_to_sim(scene_asset, scene_name, pose_tensor, env_id_tensor, env.device) def write_layout_to_sim( @@ -126,7 +107,6 @@ def write_layout_to_sim( 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) missing_assets = [ asset.name for asset in base_rotations if asset not in anchor_assets and asset not in result.positions ] @@ -142,7 +122,6 @@ def write_layout_to_sim( sim_pose, env_id, env_id_tensor, - zero_velocity, ) diff --git a/isaaclab_arena/terms/events.py b/isaaclab_arena/terms/events.py index 48409fe491..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 @@ -69,8 +70,7 @@ def _write_scene_pose(env: ManagerBasedEnv, scene_name: str, pose: Pose, env_ids 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] - asset.write_root_pose_to_sim(pose_t_xyz_q_xyzw, env_ids=env_ids) - asset.write_root_velocity_to_sim(torch.zeros(num_envs, 6, device=env.device), env_ids=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( 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" From c1f454ab78ce7dfcf31b1f870d272fa02f916dea Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Fri, 24 Jul 2026 16:31:55 +0800 Subject: [PATCH 18/19] Assert stand uses set_world_poses in placement mock tests Configure stand mocks without write_root_pose_to_sim so unit tests exercise the AssetBase dispatch path used in production for the Droid stand. Signed-off-by: Qian Lin --- isaaclab_arena/tests/test_placement_events.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index 2e769bcfb4..7176a3ee3d 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -149,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 @@ -265,6 +270,7 @@ def test_solve_and_place_objects_writes_droid_robot_and_stand(): 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 @@ -287,11 +293,12 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: ) robot_pose = env._assets["robot"].write_root_pose_to_sim.call_args.args[0] - stand_pose = env._assets["stand"].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_pose[0, 2].item() - expected_robot_z) < 1e-5 - assert abs(stand_pose[0, 0].item() - (0.2 - 0.05)) < 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(): @@ -330,6 +337,7 @@ 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) ] @@ -337,7 +345,7 @@ def test_reset_placement_asset_pose_per_env_writes_each_compound_prim_per_env(): 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"].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(): From c992a015c2b1153b33d58c8e57c74f50ecfc2427 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Fri, 24 Jul 2026 16:33:13 +0800 Subject: [PATCH 19/19] Check robot/stand orientation on kitchen reset Compare stand world quaternion to the robot root after relation placement reset so quaternion-order bugs between write_root_pose_to_sim and set_world_poses are caught, not just z alignment. Signed-off-by: Qian Lin --- isaaclab_arena/tests/test_droid_stand_height.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/isaaclab_arena/tests/test_droid_stand_height.py b/isaaclab_arena/tests/test_droid_stand_height.py index 3c8eb4d68e..d1996bd41d 100644 --- a/isaaclab_arena/tests/test_droid_stand_height.py +++ b/isaaclab_arena/tests/test_droid_stand_height.py @@ -15,6 +15,7 @@ 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: @@ -155,9 +156,21 @@ def _test_droid_kitchen_reset_alignment(simulation_app) -> bool: env = builder.make_registered() env.reset() - robot_z = wp.to_torch(env.unwrapped.scene["robot"].data.root_link_pose_w)[0, 2].item() - stand_z = env.unwrapped.scene["stand"].get_world_poses()[0][0, 2].item() + 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: