Require stable destination support for pick and place

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:

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.
 

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,
                     ),
                 ],
             ),

isaaclab_arena/tasks/predicates/spatial.py

@@ -5,17 +5,23 @@
 
 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
 
+if TYPE_CHECKING:
+    from isaaclab_arena.assets.object_base import ObjectBase
+
 
 def object_is_above_height(
     env: ManagerBasedRLEnv,
@@ -96,40 +102,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: 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."""
+
+    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: ObjectBase) -> torch.Tensor:
+    """Get the world pose of the frame in which an asset's bounding box is expressed."""
+
+    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
+
+    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: 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:
+        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: ObjectBase | None = None,
+    destination_asset: ObjectBase | 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 +259,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 +273,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