From 483f62082a1fddcc71ac2cff17e410022797d5f2 Mon Sep 17 00:00:00 2001 From: Clemens Volk Date: Wed, 29 Jul 2026 12:56:53 +0200 Subject: [PATCH 1/2] Require stable destination support for pick and place Signed-off-by: Clemens Volk --- isaaclab_arena/tasks/pick_and_place_task.py | 23 +- .../tasks/predicates/predicate_utils.py | 54 +++++ isaaclab_arena/tasks/predicates/spatial.py | 190 +++++++++++++--- .../test_object_on_destination_predicate.py | 213 ++++++++++++++++++ 4 files changed, 445 insertions(+), 35 deletions(-) create mode 100644 isaaclab_arena/tests/test_object_on_destination_predicate.py diff --git a/isaaclab_arena/tasks/pick_and_place_task.py b/isaaclab_arena/tasks/pick_and_place_task.py index bfdd9f3bf3..c17407f7db 100644 --- a/isaaclab_arena/tasks/pick_and_place_task.py +++ b/isaaclab_arena/tasks/pick_and_place_task.py @@ -35,10 +35,13 @@ @agent_ready @register_task class PickAndPlaceTask(TaskBase): - """Pick-and-place task. Success fires when the pick-up object contacts the destination - with low velocity and, when ``max_separation`` is set, is within axis-aligned proximity - of the destination. Failure (object_dropped) fires when the object falls below the - background's ``object_min_z``. + """Pick-and-place task. + + Success fires when the pick-up object rests on the destination with upward + support force and its bounding-box centroid within the destination's world + XY footprint. When ``max_separation`` is set, success also requires + axis-aligned proximity to the destination. Failure (object_dropped) fires + when the object falls below the background's ``object_min_z``. The default Mimic cfg is ``PickPlaceMimicEnvCfg``. When a task needs a different cfg shape (different arm subtask sequences, different per-subtask numerical knobs, @@ -64,6 +67,8 @@ def __init__( velocity_threshold: float = 0.1, max_separation: tuple[float, float, float] | None = None, mimic_env_cfg_factory: Callable[[ArmMode], MimicEnvCfg] | None = None, + support_cone_half_angle_deg: float = 45.0, + footprint_tolerance: float = 1e-2, ): super().__init__(episode_length_s=episode_length_s) self.pick_up_object = pick_up_object @@ -74,6 +79,8 @@ def __init__( self.scene_config = self.make_scene_cfg() self.force_threshold = force_threshold self.velocity_threshold = velocity_threshold + self.support_cone_half_angle_deg = support_cone_half_angle_deg + self.footprint_tolerance = footprint_tolerance if max_separation is not None: assert len(max_separation) == 3, f"max_separation must be (x, y, z), got {max_separation!r}" self.max_separation = max_separation @@ -113,8 +120,12 @@ def make_termination_cfg(self): params={ "object_cfg": SceneEntityCfg(self.pick_up_object.name), "contact_sensor_cfg": SceneEntityCfg(self.contact_sensor_name), + "object_asset": self.pick_up_object, + "destination_asset": self.destination_location, "force_threshold": self.force_threshold, "velocity_threshold": self.velocity_threshold, + "support_cone_half_angle_deg": self.support_cone_half_angle_deg, + "footprint_tolerance": self.footprint_tolerance, }, ), ] @@ -193,8 +204,12 @@ def get_progress_objectives(self) -> list[ProgressObjective]: object_on_destination, object_cfg=SceneEntityCfg(self.pick_up_object.name), contact_sensor_cfg=SceneEntityCfg(self.contact_sensor_name), + object_asset=self.pick_up_object, + destination_asset=self.destination_location, force_threshold=self.force_threshold, velocity_threshold=self.velocity_threshold, + support_cone_half_angle_deg=self.support_cone_half_angle_deg, + footprint_tolerance=self.footprint_tolerance, ), ], ), diff --git a/isaaclab_arena/tasks/predicates/predicate_utils.py b/isaaclab_arena/tasks/predicates/predicate_utils.py index 0942fd1357..cd22ca1b8f 100644 --- a/isaaclab_arena/tasks/predicates/predicate_utils.py +++ b/isaaclab_arena/tasks/predicates/predicate_utils.py @@ -6,9 +6,17 @@ from __future__ import annotations import torch +from typing import TYPE_CHECKING import warp as wp from isaaclab.assets import RigidObject +from isaaclab.utils.math import combine_frame_transforms + +from isaaclab_arena.assets.object_type import ObjectType +from isaaclab_arena.utils.pose import Pose + +if TYPE_CHECKING: + from isaaclab_arena.relations.placement_asset import PlaceableAsset def get_env(env): @@ -40,6 +48,52 @@ def get_root_ang_vel_w(env, name: str) -> torch.Tensor: return wp.to_torch(get_rigid_object(env, name).data.root_ang_vel_w) +def get_asset_pose_w(env, asset: PlaceableAsset) -> torch.Tensor: + """Get an asset's world pose, including unregistered object references.""" + + unwrapped_env = get_env(env) + scene_key = asset.get_scene_key() + if scene_key in unwrapped_env.scene.keys(): + if getattr(asset, "object_type", None) == ObjectType.BASE: + initial_pose = asset.get_initial_pose() + if initial_pose is None: + initial_pose = Pose.identity() + if isinstance(initial_pose, Pose): + asset_pose_w = ( + initial_pose.to_tensor(device=unwrapped_env.device).expand(unwrapped_env.num_envs, 7).clone() + ) + asset_pose_w[:, :3] += unwrapped_env.scene.env_origins + return asset_pose_w + + asset_pose_w = asset.get_object_pose(unwrapped_env, is_relative=False) + assert asset_pose_w.shape[0] in ( + 1, + unwrapped_env.num_envs, + ), f"Asset '{asset.name}' returned {asset_pose_w.shape[0]} poses for {unwrapped_env.num_envs} environments." + if asset_pose_w.shape[0] == unwrapped_env.num_envs: + return asset_pose_w + + asset_pose_w = asset_pose_w.expand(unwrapped_env.num_envs, 7).clone() + asset_pose_w[:, :3] += unwrapped_env.scene.env_origins - unwrapped_env.scene.env_origins[:1] + return asset_pose_w + + parent_asset = getattr(asset, "parent_asset", None) + pose_relative_to_parent = getattr(asset, "initial_pose_relative_to_parent", None) + assert ( + parent_asset is not None and pose_relative_to_parent is not None + ), f"Asset '{asset.name}' is not registered in the scene and has no parent-relative pose." + + parent_pose_w = get_asset_pose_w(unwrapped_env, parent_asset) + relative_pose = pose_relative_to_parent.to_tensor(device=unwrapped_env.device).expand(unwrapped_env.num_envs, 7) + position_w, quaternion_w = combine_frame_transforms( + parent_pose_w[:, :3], + parent_pose_w[:, 3:], + relative_pose[:, :3], + relative_pose[:, 3:], + ) + return torch.cat((position_w, quaternion_w), dim=-1) + + def select(result: torch.Tensor, env_id: int | None) -> torch.Tensor: """Return the entry at ``env_id`` if requested, otherwise the full vector.""" if env_id is None: diff --git a/isaaclab_arena/tasks/predicates/spatial.py b/isaaclab_arena/tasks/predicates/spatial.py index 21e3bffb06..a0a20178f3 100644 --- a/isaaclab_arena/tasks/predicates/spatial.py +++ b/isaaclab_arena/tasks/predicates/spatial.py @@ -5,16 +5,28 @@ from __future__ import annotations +import math import torch +from typing import TYPE_CHECKING import warp as wp from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv from isaaclab.managers import SceneEntityCfg from isaaclab.sensors.contact_sensor.contact_sensor import ContactSensor +from isaaclab.utils.math import combine_frame_transforms, quat_apply, quat_inv, quat_mul from isaaclab_arena.tasks.predicates.object_settling import get_object_initial_rest_state -from isaaclab_arena.tasks.predicates.predicate_utils import get_env, get_root_lin_vel_w, get_root_pos_w, select +from isaaclab_arena.tasks.predicates.predicate_utils import ( + get_asset_pose_w, + get_env, + get_root_lin_vel_w, + get_root_pos_w, + select, +) + +if TYPE_CHECKING: + from isaaclab_arena.relations.placement_asset import PlaceableAsset def object_is_above_height( @@ -96,40 +108,154 @@ def objects_in_proximity( return done +def contact_force_is_upward_support( + force_matrix_w: torch.Tensor, + force_threshold: float, + support_cone_half_angle_deg: float, +) -> torch.Tensor: + """Check whether destination contact forces support an object upward.""" + + assert ( + force_matrix_w.ndim >= 2 and force_matrix_w.shape[-1] == 3 + ), f"force_matrix_w must have shape (num_envs, ..., 3), got {tuple(force_matrix_w.shape)}" + assert force_threshold >= 0.0, f"force_threshold must be non-negative, got {force_threshold}" + assert ( + 0.0 <= support_cone_half_angle_deg < 90.0 + ), f"support_cone_half_angle_deg must be in [0, 90), got {support_cone_half_angle_deg}" + + if force_matrix_w.ndim == 2: + destination_force_w = force_matrix_w + else: + contact_axes = tuple(range(1, force_matrix_w.ndim - 1)) + destination_force_w = force_matrix_w.sum(dim=contact_axes) + force_magnitude = torch.linalg.vector_norm(destination_force_w, dim=-1) + upward_force = destination_force_w[:, 2] + minimum_upward_fraction = math.cos(math.radians(support_cone_half_angle_deg)) + + return ( + (force_magnitude >= force_threshold) + & (upward_force > 0.0) + & (upward_force >= force_magnitude * minimum_upward_fraction) + ) + + +def object_centroid_in_destination_footprint( + env: ManagerBasedRLEnv, + object_asset: PlaceableAsset, + destination_asset: PlaceableAsset, + footprint_tolerance: float = 1e-2, +) -> torch.Tensor: + """Check whether an object's bounding-box centroid is within a destination's world XY footprint.""" + + assert footprint_tolerance >= 0.0, f"footprint_tolerance must be non-negative, got {footprint_tolerance}" + + unwrapped_env = get_env(env) + object_bounding_box = _get_asset_bounding_box_per_env(object_asset, unwrapped_env.num_envs).to(unwrapped_env.device) + destination_bounding_box = _get_asset_bounding_box_per_env(destination_asset, unwrapped_env.num_envs).to( + unwrapped_env.device + ) + + object_bounding_box_pose_w = _get_bounding_box_pose_w(unwrapped_env, object_asset) + destination_bounding_box_pose_w = _get_bounding_box_pose_w(unwrapped_env, destination_asset) + object_centroid_w, _ = combine_frame_transforms( + object_bounding_box_pose_w[:, :3], + object_bounding_box_pose_w[:, 3:], + object_bounding_box.center, + ) + + destination_corners = destination_bounding_box.get_corners_at() + num_envs, num_corners, _ = destination_corners.shape + destination_quaternions = ( + destination_bounding_box_pose_w[:, None, 3:].expand(num_envs, num_corners, 4).reshape(-1, 4) + ) + destination_corners_w = quat_apply( + destination_quaternions, + destination_corners.reshape(-1, 3), + ).reshape(num_envs, num_corners, 3) + destination_corners_w += destination_bounding_box_pose_w[:, None, :3] + + minimum_xy = destination_corners_w[:, :, :2].amin(dim=1) - footprint_tolerance + maximum_xy = destination_corners_w[:, :, :2].amax(dim=1) + footprint_tolerance + return torch.all( + (object_centroid_w[:, :2] >= minimum_xy) & (object_centroid_w[:, :2] <= maximum_xy), + dim=-1, + ) + + +def _get_bounding_box_pose_w(env, asset: PlaceableAsset) -> torch.Tensor: + """Get the world pose of the frame in which an asset's bounding box is expressed.""" + + asset_pose_w = get_asset_pose_w(env, asset) + pose_relative_to_parent = getattr(asset, "initial_pose_relative_to_parent", None) + if pose_relative_to_parent is None: + return asset_pose_w + + unwrapped_env = get_env(env) + relative_pose = pose_relative_to_parent.to_tensor(device=unwrapped_env.device).expand(unwrapped_env.num_envs, 7) + bounding_box_quaternion_w = quat_mul(asset_pose_w[:, 3:], quat_inv(relative_pose[:, 3:])) + return torch.cat((asset_pose_w[:, :3], bounding_box_quaternion_w), dim=-1) + + +def _get_asset_bounding_box_per_env(asset: PlaceableAsset, num_envs: int): + """Get root-relative bounds per environment, using assigned object-set variants when available.""" + + if getattr(asset, "variant_indices_by_env", None) is not None: + return asset.get_bounding_box_per_env(num_envs) + + bounding_box = asset.get_bounding_box() + assert bounding_box.num_envs in ( + 1, + num_envs, + ), f"Asset '{asset.name}' has {bounding_box.num_envs} bounding boxes for {num_envs} environments." + if bounding_box.num_envs == num_envs: + return bounding_box + return type(bounding_box)( + min_point=bounding_box.min_point.expand(num_envs, 3), + max_point=bounding_box.max_point.expand(num_envs, 3), + ) + + def object_on_destination( env: ManagerBasedRLEnv, object_cfg: SceneEntityCfg = SceneEntityCfg("pick_up_object"), contact_sensor_cfg: SceneEntityCfg = SceneEntityCfg("pick_up_object_contact_sensor"), force_threshold: float = 1.0, velocity_threshold: float = 0.5, + object_asset: PlaceableAsset | None = None, + destination_asset: PlaceableAsset | None = None, + support_cone_half_angle_deg: float = 45.0, + footprint_tolerance: float = 1e-2, ) -> torch.Tensor: - """Checks if an object is in contact with it's destination location via a contact sensor. + """Check whether an object is resting on and within the footprint of its destination. - Returns True when the object is in contact with destination above a force threshold - and below a velocity threshold. + Returns True when destination contact supports the object upward, the object's + bounding-box centroid is within the destination's world XY footprint, and its + linear speed is below the threshold. """ unwrapped_env = get_env(env) - object: RigidObject = unwrapped_env.scene[object_cfg.name] + object_entity: RigidObject = unwrapped_env.scene[object_cfg.name] sensor: ContactSensor = unwrapped_env.scene[contact_sensor_cfg.name] + assert object_asset is not None, "object_asset is required" + assert destination_asset is not None, "destination_asset is required" + + force_matrix_w = wp.to_torch(sensor.data.force_matrix_w) + supported_by_destination = contact_force_is_upward_support( + force_matrix_w, + force_threshold=force_threshold, + support_cone_half_angle_deg=support_cone_half_angle_deg, + ) + centroid_in_footprint = object_centroid_in_destination_footprint( + env=unwrapped_env, + object_asset=object_asset, + destination_asset=destination_asset, + footprint_tolerance=footprint_tolerance, + ) - # force_matrix_w shape is (N, B, M, 3), where N is the number of sensors, B is number of bodies in each sensor - # and ``M`` is the number of filtered bodies. - # We assume B = 1 and M = 1 - assert sensor.data.force_matrix_w.shape[2] == 1 - assert sensor.data.force_matrix_w.shape[1] == 1 - # NOTE(alexmillane, 2025-08-04): We expect the binary flags to have shape (N, ) - # where N is the number of envs. - force_matrix_norm = torch.norm(wp.to_torch(sensor.data.force_matrix_w), dim=-1).reshape(-1) - force_above_threshold = force_matrix_norm > force_threshold - - velocity_w = wp.to_torch(object.data.root_lin_vel_w) - velocity_w_norm = torch.norm(velocity_w, dim=-1) - velocity_below_threshold = velocity_w_norm < velocity_threshold - - condition_met = torch.logical_and(force_above_threshold, velocity_below_threshold) + object_linear_speed = torch.linalg.vector_norm(wp.to_torch(object_entity.data.root_lin_vel_w), dim=-1) + object_at_rest = object_linear_speed < velocity_threshold - return condition_met + return supported_by_destination & centroid_in_footprint & object_at_rest def objects_on_destinations( @@ -139,10 +265,10 @@ def objects_on_destinations( force_threshold: float = 1.0, velocity_threshold: float = 0.5, ) -> torch.Tensor: - """Multi-object version of `object_on_destination`. + """Check whether every object has destination contact and low linear speed. - Returns True only when ALL objects in the list satisfy the destination condition. - See `object_on_destination` for details on the single-object logic. + This preserves the existing multi-object behavior until indirect support between + objects sharing a destination is defined. """ assert len(object_cfg_list) == len(contact_sensor_cfg_list), ( @@ -153,12 +279,14 @@ def objects_on_destinations( unwrapped_env = get_env(env) condition_met = torch.ones((unwrapped_env.num_envs), device=unwrapped_env.device, dtype=torch.bool) for object_cfg, contact_sensor_cfg in zip(object_cfg_list, contact_sensor_cfg_list): - single_condition = object_on_destination( - env=env, - object_cfg=object_cfg, - contact_sensor_cfg=contact_sensor_cfg, - force_threshold=force_threshold, - velocity_threshold=velocity_threshold, - ) + object_entity: RigidObject = unwrapped_env.scene[object_cfg.name] + sensor: ContactSensor = unwrapped_env.scene[contact_sensor_cfg.name] + assert sensor.data.force_matrix_w.shape[2] == 1 + assert sensor.data.force_matrix_w.shape[1] == 1 + + force_matrix_norm = torch.linalg.vector_norm(wp.to_torch(sensor.data.force_matrix_w), dim=-1).reshape(-1) + force_above_threshold = force_matrix_norm > force_threshold + object_linear_speed = torch.linalg.vector_norm(wp.to_torch(object_entity.data.root_lin_vel_w), dim=-1) + single_condition = force_above_threshold & (object_linear_speed < velocity_threshold) condition_met = torch.logical_and(condition_met, single_condition) return condition_met diff --git a/isaaclab_arena/tests/test_object_on_destination_predicate.py b/isaaclab_arena/tests/test_object_on_destination_predicate.py new file mode 100644 index 0000000000..ac4b58e501 --- /dev/null +++ b/isaaclab_arena/tests/test_object_on_destination_predicate.py @@ -0,0 +1,213 @@ +# 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 math +import torch + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + + +def _test_object_on_destination_predicate(_simulation_app) -> bool: + import warp as wp + from isaaclab.managers import SceneEntityCfg + + from isaaclab_arena.tasks.predicates.spatial import ( + contact_force_is_upward_support, + object_centroid_in_destination_footprint, + object_on_destination, + ) + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.pose import Pose + + class DummyScene(dict): + def __init__(self, num_envs: int): + super().__init__() + self.env_origins = torch.zeros((num_envs, 3)) + + class DummyEnv: + def __init__(self, num_envs: int): + self.num_envs = num_envs + self.device = "cpu" + self.scene = DummyScene(num_envs) + + class DummyAsset: + def __init__(self, name: str, bounding_box: AxisAlignedBoundingBox, pose_w: torch.Tensor): + self.name = name + self._bounding_box = bounding_box + self._pose_w = pose_w + + def get_scene_key(self) -> str: + return self.name + + def get_bounding_box(self) -> AxisAlignedBoundingBox: + return self._bounding_box + + def get_object_pose(self, _env, is_relative: bool = True) -> torch.Tensor: + assert not is_relative + return self._pose_w + + identity_quaternion = (0.0, 0.0, 0.0, 1.0) + + force_matrix_w = torch.tensor([ + [[[0.0, 0.0, 0.05], [0.0, 0.0, 0.0]]], + [[[0.0, 0.0, 0.2], [0.0, 0.0, 0.0]]], + [[[0.2, 0.0, 0.0], [0.0, 0.0, 0.0]]], + [[[0.0, 0.0, -0.2], [0.0, 0.0, 0.0]]], + [[[1.0, 0.0, 1.0], [0.0, 0.0, 0.0]]], + [[[1.01, 0.0, 1.0], [0.0, 0.0, 0.0]]], + [[[1.0, 0.0, 0.5], [-1.0, 0.0, 0.5]]], + ]) + support_result = contact_force_is_upward_support( + force_matrix_w, + force_threshold=0.1, + support_cone_half_angle_deg=45.0, + ) + torch.testing.assert_close( + support_result, + torch.tensor([False, True, False, False, True, False, True]), + ) + torch.testing.assert_close( + contact_force_is_upward_support( + torch.tensor([[0.0, 0.0, 0.2], [0.2, 0.0, 0.0]]), + force_threshold=0.1, + support_cone_half_angle_deg=45.0, + ), + torch.tensor([True, False]), + ) + + num_envs = 5 + env = DummyEnv(num_envs) + yaw_90_quaternion = (0.0, 0.0, math.sqrt(0.5), math.sqrt(0.5)) + object_pose_w = torch.tensor([ + [-0.2, 0.0, 0.0, *identity_quaternion], + [0.8, 0.5, 0.0, *identity_quaternion], + # The object overlaps the destination from outside, but its centroid + # remains beyond the destination footprint. + [0.81, 0.0, 0.0, *identity_quaternion], + [0.2, 0.9, 100.0, *identity_quaternion], + [0.4, 0.0, 0.0, *identity_quaternion], + ]) + destination_pose_w = torch.tensor([ + [0.0, 0.0, 0.0, *identity_quaternion], + [0.0, 0.0, 0.0, *identity_quaternion], + [0.0, 0.0, 0.0, *identity_quaternion], + [0.0, 0.0, 0.0, *yaw_90_quaternion], + [0.0, 0.0, 0.0, *yaw_90_quaternion], + ]) + object_asset = DummyAsset( + "object", + AxisAlignedBoundingBox(min_point=(0.1, -0.1, -0.1), max_point=(0.3, 0.1, 0.1)), + object_pose_w, + ) + destination_asset = DummyAsset( + "destination", + AxisAlignedBoundingBox(min_point=(-1.0, -0.5, -0.2), max_point=(1.0, 0.5, 0.2)), + destination_pose_w, + ) + env.scene[object_asset.name] = object() + env.scene[destination_asset.name] = object() + + footprint_result = object_centroid_in_destination_footprint( + env, + object_asset=object_asset, + destination_asset=destination_asset, + footprint_tolerance=0.0, + ) + torch.testing.assert_close( + footprint_result, + torch.tensor([True, True, False, True, False]), + ) + + parent_pose_w = torch.tensor([[1.0, 2.0, 0.0, *yaw_90_quaternion]]) + reference_env = DummyEnv(num_envs=1) + parent_asset = DummyAsset( + "parent", + AxisAlignedBoundingBox(min_point=(-2.0, -2.0, -0.2), max_point=(2.0, 2.0, 0.2)), + parent_pose_w, + ) + reference_env.scene[parent_asset.name] = object() + reference_object_asset = DummyAsset( + "reference_object", + AxisAlignedBoundingBox(min_point=(-0.1, -0.1, -0.1), max_point=(0.1, 0.1, 0.1)), + torch.tensor([[1.25, 1.75, 0.0, *identity_quaternion]]), + ) + reference_env.scene[reference_object_asset.name] = object() + + reference_destination_asset = DummyAsset( + "reference_destination", + AxisAlignedBoundingBox(min_point=(-1.0, -0.5, -0.2), max_point=(1.0, 0.5, 0.2)), + torch.empty((1, 7)), + ) + reference_destination_asset.parent_asset = parent_asset + reference_destination_asset.initial_pose_relative_to_parent = Pose( + position_xyz=(0.5, 0.0, 0.0), + rotation_xyzw=yaw_90_quaternion, + ) + reference_footprint_result = object_centroid_in_destination_footprint( + reference_env, + object_asset=reference_object_asset, + destination_asset=reference_destination_asset, + footprint_tolerance=0.0, + ) + assert reference_footprint_result.item() + + combined_env = DummyEnv(num_envs=3) + combined_object_asset = DummyAsset( + "object", + AxisAlignedBoundingBox(min_point=(0.1, -0.1, -0.1), max_point=(0.3, 0.1, 0.1)), + torch.tensor([ + [-0.2, 0.0, 0.0, *identity_quaternion], + [-0.2, 0.0, 0.0, *identity_quaternion], + [0.81, 0.0, 0.0, *identity_quaternion], + ]), + ) + combined_destination_asset = DummyAsset( + "destination", + AxisAlignedBoundingBox(min_point=(-1.0, -0.5, -0.2), max_point=(1.0, 0.5, 0.2)), + torch.tensor([[0.0, 0.0, 0.0, *identity_quaternion]]).expand(3, 7), + ) + + class DummyRigidObject: + def __init__(self): + self.data = type( + "DummyRigidObjectData", + (), + {"root_lin_vel_w": wp.from_torch(torch.zeros((3, 3)), dtype=wp.vec3)}, + )() + + class DummyContactSensor: + def __init__(self): + force_matrix_w = torch.tensor([ + [[[0.0, 0.0, 0.2]]], + [[[0.2, 0.0, 0.0]]], + [[[0.0, 0.0, 0.2]]], + ]) + self.data = type( + "DummyContactSensorData", + (), + {"force_matrix_w": wp.from_torch(force_matrix_w, dtype=wp.vec3)}, + )() + + combined_env.scene[combined_object_asset.name] = DummyRigidObject() + combined_env.scene[combined_destination_asset.name] = object() + combined_env.scene["contact_sensor"] = DummyContactSensor() + combined_result = object_on_destination( + combined_env, + object_cfg=SceneEntityCfg(combined_object_asset.name), + contact_sensor_cfg=SceneEntityCfg("contact_sensor"), + object_asset=combined_object_asset, + destination_asset=combined_destination_asset, + force_threshold=0.1, + velocity_threshold=0.1, + support_cone_half_angle_deg=45.0, + footprint_tolerance=0.0, + ) + torch.testing.assert_close(combined_result, torch.tensor([True, False, False])) + + return True + + +def test_object_on_destination_predicate(): + assert run_simulation_app_function(_test_object_on_destination_predicate) From 3de99f92a8e1bab4cf532d8265ffb838479c6d87 Mon Sep 17 00:00:00 2001 From: Clemens Volk Date: Wed, 29 Jul 2026 14:19:33 +0200 Subject: [PATCH 2/2] Centralize runtime pose lookup Signed-off-by: Clemens Volk --- isaaclab_arena/assets/object_base.py | 17 ++-- isaaclab_arena/assets/object_reference.py | 26 ++++++ .../tasks/predicates/predicate_utils.py | 54 ------------- isaaclab_arena/tasks/predicates/spatial.py | 24 +++--- .../tests/test_object_of_type_base.py | 17 +++- .../test_object_on_destination_predicate.py | 79 ++++++++++++++++--- 6 files changed, 128 insertions(+), 89 deletions(-) diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py index ba7635310e..134d1174a5 100644 --- a/isaaclab_arena/assets/object_base.py +++ b/isaaclab_arena/assets/object_base.py @@ -157,16 +157,23 @@ def get_object_pose(self, env: ManagerBasedEnv, is_relative: bool = True) -> tor The pose of the object in each environment. The shape is (num_envs, 7). The order is (x, y, z, qx, qy, qz, qw). """ - # We require that the asset has been added to the scene under its name. - assert self.name in env.unwrapped.scene.keys(), f"Asset {self.name} not found in scene" + unwrapped_env = env.unwrapped + scene_key = self.get_scene_key() + assert scene_key in unwrapped_env.scene.keys(), f"Asset {self.name} not found in scene" if (self.object_type == ObjectType.RIGID) or (self.object_type == ObjectType.ARTICULATION): - object_pose = wp.to_torch(env.unwrapped.scene[self.name].data.root_pose_w).clone() + object_pose = wp.to_torch(unwrapped_env.scene[scene_key].data.root_pose_w).clone() elif self.object_type == ObjectType.BASE: - object_pose = torch.cat(env.unwrapped.scene[self.name].get_world_poses(), dim=-1) + initial_pose = self._get_initial_pose_as_pose() or Pose.identity() + object_pose = initial_pose.to_tensor(device=unwrapped_env.device).to( + dtype=unwrapped_env.scene.env_origins.dtype + ) + object_pose = object_pose.expand(unwrapped_env.num_envs, -1).clone() + object_pose[:, :3] += unwrapped_env.scene.env_origins else: raise ValueError(f"Function not implemented for object type: {self.object_type}") + if is_relative: - object_pose[:, :3] -= env.unwrapped.scene.env_origins + object_pose[:, :3] -= unwrapped_env.scene.env_origins return object_pose def set_object_pose(self, env: ManagerBasedEnv, pose: Pose, env_ids: torch.Tensor | None = None) -> None: diff --git a/isaaclab_arena/assets/object_reference.py b/isaaclab_arena/assets/object_reference.py index 9310b07c4d..c9356a956b 100644 --- a/isaaclab_arena/assets/object_reference.py +++ b/isaaclab_arena/assets/object_reference.py @@ -3,10 +3,13 @@ # # SPDX-License-Identifier: Apache-2.0 +import torch import trimesh from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg +from isaaclab.envs import ManagerBasedEnv from isaaclab.sensors.contact_sensor.contact_sensor_cfg import ContactSensorCfg +from isaaclab.utils.math import combine_frame_transforms from pxr import Usd from isaaclab_arena.affordances.openable import Openable @@ -48,6 +51,29 @@ def get_initial_pose(self) -> Pose: T_W_O = T_W_P.multiply(T_P_O) return T_W_O + def get_object_pose(self, env: ManagerBasedEnv, is_relative: bool = True) -> torch.Tensor: + """Get the reference pose from its scene entity or parent.""" + + unwrapped_env = env.unwrapped + if self.get_scene_key() in unwrapped_env.scene.keys(): + return super().get_object_pose(env, is_relative=is_relative) + + parent_pose_w = self.parent_asset.get_object_pose(unwrapped_env, is_relative=False) + relative_pose = self.initial_pose_relative_to_parent.to_tensor(device=parent_pose_w.device).to( + dtype=parent_pose_w.dtype + ) + relative_pose = relative_pose.expand(parent_pose_w.shape[0], -1) + position_w, quaternion_w = combine_frame_transforms( + parent_pose_w[:, :3], + parent_pose_w[:, 3:], + relative_pose[:, :3], + relative_pose[:, 3:], + ) + object_pose = torch.cat((position_w, quaternion_w), dim=-1) + if is_relative: + object_pose[:, :3] -= unwrapped_env.scene.env_origins + return object_pose + def add_relation(self, relation: RelationBase) -> None: """Add a relation to this object reference. diff --git a/isaaclab_arena/tasks/predicates/predicate_utils.py b/isaaclab_arena/tasks/predicates/predicate_utils.py index cd22ca1b8f..0942fd1357 100644 --- a/isaaclab_arena/tasks/predicates/predicate_utils.py +++ b/isaaclab_arena/tasks/predicates/predicate_utils.py @@ -6,17 +6,9 @@ from __future__ import annotations import torch -from typing import TYPE_CHECKING import warp as wp from isaaclab.assets import RigidObject -from isaaclab.utils.math import combine_frame_transforms - -from isaaclab_arena.assets.object_type import ObjectType -from isaaclab_arena.utils.pose import Pose - -if TYPE_CHECKING: - from isaaclab_arena.relations.placement_asset import PlaceableAsset def get_env(env): @@ -48,52 +40,6 @@ def get_root_ang_vel_w(env, name: str) -> torch.Tensor: return wp.to_torch(get_rigid_object(env, name).data.root_ang_vel_w) -def get_asset_pose_w(env, asset: PlaceableAsset) -> torch.Tensor: - """Get an asset's world pose, including unregistered object references.""" - - unwrapped_env = get_env(env) - scene_key = asset.get_scene_key() - if scene_key in unwrapped_env.scene.keys(): - if getattr(asset, "object_type", None) == ObjectType.BASE: - initial_pose = asset.get_initial_pose() - if initial_pose is None: - initial_pose = Pose.identity() - if isinstance(initial_pose, Pose): - asset_pose_w = ( - initial_pose.to_tensor(device=unwrapped_env.device).expand(unwrapped_env.num_envs, 7).clone() - ) - asset_pose_w[:, :3] += unwrapped_env.scene.env_origins - return asset_pose_w - - asset_pose_w = asset.get_object_pose(unwrapped_env, is_relative=False) - assert asset_pose_w.shape[0] in ( - 1, - unwrapped_env.num_envs, - ), f"Asset '{asset.name}' returned {asset_pose_w.shape[0]} poses for {unwrapped_env.num_envs} environments." - if asset_pose_w.shape[0] == unwrapped_env.num_envs: - return asset_pose_w - - asset_pose_w = asset_pose_w.expand(unwrapped_env.num_envs, 7).clone() - asset_pose_w[:, :3] += unwrapped_env.scene.env_origins - unwrapped_env.scene.env_origins[:1] - return asset_pose_w - - parent_asset = getattr(asset, "parent_asset", None) - pose_relative_to_parent = getattr(asset, "initial_pose_relative_to_parent", None) - assert ( - parent_asset is not None and pose_relative_to_parent is not None - ), f"Asset '{asset.name}' is not registered in the scene and has no parent-relative pose." - - parent_pose_w = get_asset_pose_w(unwrapped_env, parent_asset) - relative_pose = pose_relative_to_parent.to_tensor(device=unwrapped_env.device).expand(unwrapped_env.num_envs, 7) - position_w, quaternion_w = combine_frame_transforms( - parent_pose_w[:, :3], - parent_pose_w[:, 3:], - relative_pose[:, :3], - relative_pose[:, 3:], - ) - return torch.cat((position_w, quaternion_w), dim=-1) - - def select(result: torch.Tensor, env_id: int | None) -> torch.Tensor: """Return the entry at ``env_id`` if requested, otherwise the full vector.""" if env_id is None: diff --git a/isaaclab_arena/tasks/predicates/spatial.py b/isaaclab_arena/tasks/predicates/spatial.py index a0a20178f3..51cb0d35ff 100644 --- a/isaaclab_arena/tasks/predicates/spatial.py +++ b/isaaclab_arena/tasks/predicates/spatial.py @@ -17,16 +17,10 @@ from isaaclab.utils.math import combine_frame_transforms, quat_apply, quat_inv, quat_mul from isaaclab_arena.tasks.predicates.object_settling import get_object_initial_rest_state -from isaaclab_arena.tasks.predicates.predicate_utils import ( - get_asset_pose_w, - get_env, - get_root_lin_vel_w, - get_root_pos_w, - select, -) +from isaaclab_arena.tasks.predicates.predicate_utils import get_env, get_root_lin_vel_w, get_root_pos_w, select if TYPE_CHECKING: - from isaaclab_arena.relations.placement_asset import PlaceableAsset + from isaaclab_arena.assets.object_base import ObjectBase def object_is_above_height( @@ -141,8 +135,8 @@ def contact_force_is_upward_support( def object_centroid_in_destination_footprint( env: ManagerBasedRLEnv, - object_asset: PlaceableAsset, - destination_asset: PlaceableAsset, + object_asset: ObjectBase, + destination_asset: ObjectBase, footprint_tolerance: float = 1e-2, ) -> torch.Tensor: """Check whether an object's bounding-box centroid is within a destination's world XY footprint.""" @@ -182,10 +176,10 @@ def object_centroid_in_destination_footprint( ) -def _get_bounding_box_pose_w(env, asset: PlaceableAsset) -> torch.Tensor: +def _get_bounding_box_pose_w(env, asset: ObjectBase) -> torch.Tensor: """Get the world pose of the frame in which an asset's bounding box is expressed.""" - asset_pose_w = get_asset_pose_w(env, asset) + asset_pose_w = asset.get_object_pose(env, is_relative=False) pose_relative_to_parent = getattr(asset, "initial_pose_relative_to_parent", None) if pose_relative_to_parent is None: return asset_pose_w @@ -196,7 +190,7 @@ def _get_bounding_box_pose_w(env, asset: PlaceableAsset) -> torch.Tensor: return torch.cat((asset_pose_w[:, :3], bounding_box_quaternion_w), dim=-1) -def _get_asset_bounding_box_per_env(asset: PlaceableAsset, num_envs: int): +def _get_asset_bounding_box_per_env(asset: ObjectBase, num_envs: int): """Get root-relative bounds per environment, using assigned object-set variants when available.""" if getattr(asset, "variant_indices_by_env", None) is not None: @@ -221,8 +215,8 @@ def object_on_destination( contact_sensor_cfg: SceneEntityCfg = SceneEntityCfg("pick_up_object_contact_sensor"), force_threshold: float = 1.0, velocity_threshold: float = 0.5, - object_asset: PlaceableAsset | None = None, - destination_asset: PlaceableAsset | None = None, + object_asset: ObjectBase | None = None, + destination_asset: ObjectBase | None = None, support_cone_half_angle_deg: float = 45.0, footprint_tolerance: float = 1e-2, ) -> torch.Tensor: diff --git a/isaaclab_arena/tests/test_object_of_type_base.py b/isaaclab_arena/tests/test_object_of_type_base.py index efbff233f3..0d71b86f48 100644 --- a/isaaclab_arena/tests/test_object_of_type_base.py +++ b/isaaclab_arena/tests/test_object_of_type_base.py @@ -59,11 +59,24 @@ def __init__(self, prim_path: str = default_prim_path, initial_pose: Pose | None try: args_cli = get_isaaclab_arena_cli_parser().parse_args([]) + args_cli.num_envs = 2 env_builder = ArenaEnvBuilder(isaaclab_arena_environment, arena_env_builder_cfg_from_argparse(args_cli)) env = env_builder.make_registered() env.reset() - position_before_simulation = torch.tensor(cone.get_initial_pose().position_xyz) + expected_relative_pose = ( + cone.get_initial_pose().to_tensor(device=env.unwrapped.device).expand(env.unwrapped.num_envs, 7) + ) + actual_relative_pose = cone.get_object_pose(env, is_relative=True) + torch.testing.assert_close(actual_relative_pose, expected_relative_pose) + + expected_world_pose = expected_relative_pose.clone() + expected_world_pose[:, :3] += env.unwrapped.scene.env_origins + actual_world_pose = cone.get_object_pose(env, is_relative=False) + torch.testing.assert_close(actual_world_pose, expected_world_pose) + + position_before_simulation, _ = env.unwrapped.scene["cone_no_physics"].get_world_poses() + position_before_simulation = position_before_simulation.torch.clone().cpu() # Run some zero actions. for _ in tqdm.tqdm(range(NUM_STEPS)): @@ -73,7 +86,7 @@ def __init__(self, prim_path: str = default_prim_path, initial_pose: Pose | None # Check the the object is floating. position_after_simulation, _ = env.unwrapped.scene["cone_no_physics"].get_world_poses() - movement = position_after_simulation.cpu() - position_before_simulation.cpu() + movement = position_after_simulation.torch.cpu() - position_before_simulation assert torch.norm(movement).item() < MOVEMENT_EPS, "Object moved. Should not have physics." except Exception as e: diff --git a/isaaclab_arena/tests/test_object_on_destination_predicate.py b/isaaclab_arena/tests/test_object_on_destination_predicate.py index ac4b58e501..050432855a 100644 --- a/isaaclab_arena/tests/test_object_on_destination_predicate.py +++ b/isaaclab_arena/tests/test_object_on_destination_predicate.py @@ -13,6 +13,9 @@ def _test_object_on_destination_predicate(_simulation_app) -> bool: import warp as wp from isaaclab.managers import SceneEntityCfg + from isaaclab_arena.assets.object import Object + from isaaclab_arena.assets.object_base import ObjectType + from isaaclab_arena.assets.object_reference import ObjectReference from isaaclab_arena.tasks.predicates.spatial import ( contact_force_is_upward_support, object_centroid_in_destination_footprint, @@ -31,6 +34,7 @@ def __init__(self, num_envs: int): self.num_envs = num_envs self.device = "cpu" self.scene = DummyScene(num_envs) + self.unwrapped = self class DummyAsset: def __init__(self, name: str, bounding_box: AxisAlignedBoundingBox, pose_w: torch.Tensor): @@ -120,38 +124,87 @@ def get_object_pose(self, _env, is_relative: bool = True) -> torch.Tensor: torch.tensor([True, True, False, True, False]), ) - parent_pose_w = torch.tensor([[1.0, 2.0, 0.0, *yaw_90_quaternion]]) - reference_env = DummyEnv(num_envs=1) - parent_asset = DummyAsset( - "parent", - AxisAlignedBoundingBox(min_point=(-2.0, -2.0, -0.2), max_point=(2.0, 2.0, 0.2)), - parent_pose_w, + reference_env = DummyEnv(num_envs=2) + reference_env.scene.env_origins = torch.tensor([ + [0.0, 0.0, 0.0], + [4.0, -3.0, 1.0], + ]) + parent_asset = Object.__new__(Object) + parent_asset.name = "parent" + parent_asset.object_type = ObjectType.BASE + parent_asset.initial_pose = Pose( + position_xyz=(1.0, 2.0, 0.0), + rotation_xyzw=yaw_90_quaternion, ) reference_env.scene[parent_asset.name] = object() + + expected_parent_pose_relative = torch.tensor([ + [1.0, 2.0, 0.0, *yaw_90_quaternion], + [1.0, 2.0, 0.0, *yaw_90_quaternion], + ]) + expected_parent_pose_w = expected_parent_pose_relative.clone() + expected_parent_pose_w[:, :3] += reference_env.scene.env_origins + torch.testing.assert_close(parent_asset.get_object_pose(reference_env), expected_parent_pose_relative) + torch.testing.assert_close( + parent_asset.get_object_pose(reference_env, is_relative=False), + expected_parent_pose_w, + ) + + parent_without_pose = Object.__new__(Object) + parent_without_pose.name = "parent_without_pose" + parent_without_pose.object_type = ObjectType.BASE + parent_without_pose.initial_pose = None + reference_env.scene[parent_without_pose.name] = object() + expected_identity_pose = torch.tensor([ + [0.0, 0.0, 0.0, *identity_quaternion], + [0.0, 0.0, 0.0, *identity_quaternion], + ]) + torch.testing.assert_close(parent_without_pose.get_object_pose(reference_env), expected_identity_pose) + reference_object_asset = DummyAsset( "reference_object", AxisAlignedBoundingBox(min_point=(-0.1, -0.1, -0.1), max_point=(0.1, 0.1, 0.1)), - torch.tensor([[1.25, 1.75, 0.0, *identity_quaternion]]), + torch.tensor([ + [1.25, 1.75, 0.0, *identity_quaternion], + [5.25, -1.25, 1.0, *identity_quaternion], + ]), ) reference_env.scene[reference_object_asset.name] = object() - reference_destination_asset = DummyAsset( - "reference_destination", - AxisAlignedBoundingBox(min_point=(-1.0, -0.5, -0.2), max_point=(1.0, 0.5, 0.2)), - torch.empty((1, 7)), - ) + reference_destination_asset = ObjectReference.__new__(ObjectReference) + reference_destination_asset.name = "reference_destination" reference_destination_asset.parent_asset = parent_asset reference_destination_asset.initial_pose_relative_to_parent = Pose( position_xyz=(0.5, 0.0, 0.0), rotation_xyzw=yaw_90_quaternion, ) + reference_destination_asset._bounding_box = AxisAlignedBoundingBox( + min_point=(-1.0, -0.5, -0.2), + max_point=(1.0, 0.5, 0.2), + ) + + expected_reference_pose_relative = torch.tensor([ + [1.0, 2.5, 0.0, 0.0, 0.0, 1.0, 0.0], + [1.0, 2.5, 0.0, 0.0, 0.0, 1.0, 0.0], + ]) + expected_reference_pose_w = expected_reference_pose_relative.clone() + expected_reference_pose_w[:, :3] += reference_env.scene.env_origins + torch.testing.assert_close( + reference_destination_asset.get_object_pose(reference_env), + expected_reference_pose_relative, + ) + torch.testing.assert_close( + reference_destination_asset.get_object_pose(reference_env, is_relative=False), + expected_reference_pose_w, + ) + reference_footprint_result = object_centroid_in_destination_footprint( reference_env, object_asset=reference_object_asset, destination_asset=reference_destination_asset, footprint_tolerance=0.0, ) - assert reference_footprint_result.item() + torch.testing.assert_close(reference_footprint_result, torch.tensor([True, True])) combined_env = DummyEnv(num_envs=3) combined_object_asset = DummyAsset(