return tensor for all bbox

Author: zhx06

isaaclab_arena/assets/object_set.py

@@ -88,7 +88,7 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox:
         Returns the bounding box with the greatest z-extent among all objects in the set.
         This is a heuristic to avoid objects spawning inside their support surfaces.
         """
-        return max(self.objects, key=lambda obj: obj.get_bounding_box().size[2]).get_bounding_box()
+        return max(self.objects, key=lambda obj: obj.get_bounding_box().size[0, 2].item()).get_bounding_box()
 
     def get_contact_sensor_cfg(self, contact_against_prim_paths: list[str] | None = None) -> ContactSensorCfg:
         # We override this function from the parent class because in some assets, the rigid body

isaaclab_arena/relations/object_placer.py

@@ -143,24 +143,12 @@ def _get_init_bounds(self, anchor_object: Object | ObjectReference) -> AxisAlign
 
         # Create bounds centered on anchor's world bounding box center
         anchor_bbox = anchor_object.get_world_bounding_box()
-        center = anchor_bbox.center
-        half_size = (
-            self.params.init_bounds_size[0] / 2,
-            self.params.init_bounds_size[1] / 2,
-            self.params.init_bounds_size[2] / 2,
-        )
+        center = anchor_bbox.center[0]  # (3,) tensor
+        half_size = torch.tensor(self.params.init_bounds_size) / 2
 
         return AxisAlignedBoundingBox(
-            min_point=(
-                center[0] - half_size[0],
-                center[1] - half_size[1],
-                center[2] - half_size[2],
-            ),
-            max_point=(
-                center[0] + half_size[0],
-                center[1] + half_size[1],
-                center[2] + half_size[2],
-            ),
+            min_point=center - half_size,
+            max_point=center + half_size,
         )
 
     def _generate_initial_positions(
@@ -206,17 +194,17 @@ def _validate_on_relations(
                 parent_world = parent.get_bounding_box().translated(positions[parent])
                 # 1 & 2: Same as OnLossStrategy X/Y band (child's footprint within parent).
                 if (
-                    child_world.min_point[0] < parent_world.min_point[0]
-                    or child_world.max_point[0] > parent_world.max_point[0]
-                    or child_world.min_point[1] < parent_world.min_point[1]
-                    or child_world.max_point[1] > parent_world.max_point[1]
+                    child_world.min_point[0, 0] < parent_world.min_point[0, 0]
+                    or child_world.max_point[0, 0] > parent_world.max_point[0, 0]
+                    or child_world.min_point[0, 1] < parent_world.min_point[0, 1]
+                    or child_world.max_point[0, 1] > parent_world.max_point[0, 1]
                 ):
                     if self.params.verbose:
                         print(f"  On relation: '{obj.name}' XY outside parent (retrying)")
                     return False
                 # 3. Z: same as OnLossStrategy; child_bottom in (parent_top, parent_top+clearance_m], within on_relation_z_tolerance_m.
-                parent_local_top_z: float = parent.get_bounding_box().max_point[2]
-                child_local_bottom_z: float = obj.get_bounding_box().min_point[2]
+                parent_local_top_z: float = parent.get_bounding_box().max_point[0, 2].item()
+                child_local_bottom_z: float = obj.get_bounding_box().min_point[0, 2].item()
                 parent_top_z = parent_local_top_z + positions[parent][2]
                 clearance_m = rel.clearance_m
                 child_bottom_z = child_local_bottom_z + positions[obj][2]
@@ -240,7 +228,7 @@ def _validate_no_overlap(
                 a_world = a.get_bounding_box().translated(positions[a])
                 b_world = b.get_bounding_box().translated(positions[b])
 
-                if a_world.overlaps(b_world, margin=self.params.min_separation_m):
+                if a_world.overlaps(b_world, margin=self.params.min_separation_m).item():
                     if self.params.verbose:
                         print(f"  Overlap between '{a.name}' and '{b.name}'")
                     return False

isaaclab_arena/relations/relation_loss_strategies.py

@@ -158,12 +158,12 @@ def compute_loss(
 
         # Parent world extents from the world bounding box
         if cfg.direction == Direction.POSITIVE:
-            parent_edge = parent_world_bbox.max_point[cfg.primary_axis]
-            child_offset = child_bbox.min_point[cfg.primary_axis]
+            parent_edge = parent_world_bbox.max_point[0, cfg.primary_axis]
+            child_offset = child_bbox.min_point[0, cfg.primary_axis]
             penalty_side = "less"
         else:
-            parent_edge = parent_world_bbox.min_point[cfg.primary_axis]
-            child_offset = child_bbox.max_point[cfg.primary_axis]
+            parent_edge = parent_world_bbox.min_point[0, cfg.primary_axis]
+            child_offset = child_bbox.max_point[0, cfg.primary_axis]
             penalty_side = "greater"
 
         # 1. Half-plane loss: child must be on correct side of parent edge
@@ -175,10 +175,10 @@ def compute_loss(
         )
 
         # 2. Band position loss: child placed at target position within parent's perpendicular extent
-        parent_band_min = parent_world_bbox.min_point[cfg.band_axis]
-        parent_band_max = parent_world_bbox.max_point[cfg.band_axis]
-        valid_band_min = parent_band_min - child_bbox.min_point[cfg.band_axis]
-        valid_band_max = parent_band_max - child_bbox.max_point[cfg.band_axis]
+        parent_band_min = parent_world_bbox.min_point[0, cfg.band_axis]
+        parent_band_max = parent_world_bbox.max_point[0, cfg.band_axis]
+        valid_band_min = parent_band_min - child_bbox.min_point[0, cfg.band_axis]
+        valid_band_max = parent_band_max - child_bbox.max_point[0, cfg.band_axis]
         # Convert cross_position_ratio [-1, 1] to interpolation factor [0, 1]: -1 = min, 0 = center, 1 = max
         t = (relation.cross_position_ratio + 1.0) / 2.0
         target_band_pos = valid_band_min + t * (valid_band_max - valid_band_min)
@@ -199,19 +199,19 @@ def compute_loss(
             band_axis_name = cfg.band_axis.name
             print(
                 f"    [NextTo] {relation.side.value}: child_{axis_name.lower()}="
-                f"{child_pos[cfg.primary_axis].item():.4f}, parent_edge={parent_edge:.4f},"
+                f"{child_pos[cfg.primary_axis].item():.4f}, parent_edge={parent_edge.item():.4f},"
                 f" loss={half_plane_loss.item():.6f}"
             )
             print(
                 f"    [NextTo] {band_axis_name} band: child_{band_axis_name.lower()}="
-                f"{child_pos[cfg.band_axis].item():.4f}, target={target_band_pos:.4f}"
+                f"{child_pos[cfg.band_axis].item():.4f}, target={target_band_pos.item():.4f}"
                 f" (cross_position_ratio={relation.cross_position_ratio:.2f},"
-                f" range=[{valid_band_min:.4f}, {valid_band_max:.4f}]),"
+                f" range=[{valid_band_min.item():.4f}, {valid_band_max.item():.4f}]),"
                 f" loss={band_loss.item():.6f}"
             )
             print(
                 f"    [NextTo] Distance: child_{axis_name.lower()}="
-                f"{child_pos[cfg.primary_axis].item():.4f}, target={target_pos:.4f},"
+                f"{child_pos[cfg.primary_axis].item():.4f}, target={target_pos.item():.4f},"
                 f" loss={distance_loss.item():.6f}"
             )
 
@@ -257,19 +257,19 @@ def compute_loss(
             Weighted loss tensor.
         """
         # Parent world-space extents from the world bounding box
-        parent_x_min = parent_world_bbox.min_point[0]
-        parent_x_max = parent_world_bbox.max_point[0]
-        parent_y_min = parent_world_bbox.min_point[1]
-        parent_y_max = parent_world_bbox.max_point[1]
-        parent_z_max = parent_world_bbox.max_point[2]  # Top surface
+        parent_x_min = parent_world_bbox.min_point[0, 0]
+        parent_x_max = parent_world_bbox.max_point[0, 0]
+        parent_y_min = parent_world_bbox.min_point[0, 1]
+        parent_y_max = parent_world_bbox.max_point[0, 1]
+        parent_z_max = parent_world_bbox.max_point[0, 2]  # Top surface
 
         # Compute valid position ranges such that child's entire footprint is within parent
-        # Child left edge = child_pos[0] + child_bbox.min_point[0], must be >= parent_x_min
-        # Child right edge = child_pos[0] + child_bbox.max_point[0], must be <= parent_x_max
-        valid_x_min = parent_x_min - child_bbox.min_point[0]  # child's left at parent's left
-        valid_x_max = parent_x_max - child_bbox.max_point[0]  # child's right at parent's right
-        valid_y_min = parent_y_min - child_bbox.min_point[1]
-        valid_y_max = parent_y_max - child_bbox.max_point[1]
+        # Child left edge = child_pos[0] + child_bbox.min_point[0, 0], must be >= parent_x_min
+        # Child right edge = child_pos[0] + child_bbox.max_point[0, 0], must be <= parent_x_max
+        valid_x_min = parent_x_min - child_bbox.min_point[0, 0]  # child's left at parent's left
+        valid_x_max = parent_x_max - child_bbox.max_point[0, 0]  # child's right at parent's right
+        valid_y_min = parent_y_min - child_bbox.min_point[0, 1]
+        valid_y_max = parent_y_max - child_bbox.max_point[0, 1]
 
         # 1. X band loss: child's footprint entirely within parent's X extent
         x_band_loss = linear_band_loss(
@@ -288,19 +288,22 @@ def compute_loss(
         )
 
         # 3. Z point loss: child bottom = parent top + clearance
-        target_z = parent_z_max + relation.clearance_m - child_bbox.min_point[2]
+        target_z = parent_z_max + relation.clearance_m - child_bbox.min_point[0, 2]
         z_loss = single_point_linear_loss(child_pos[2], target_z, slope=self.slope)
 
         if self.debug:
             print(
-                f"    [On] X: child_pos={child_pos[0].item():.4f}, valid_range=[{valid_x_min:.4f},"
-                f" {valid_x_max:.4f}], loss={x_band_loss.item():.6f}"
+                f"    [On] X: child_pos={child_pos[0].item():.4f}, valid_range=[{valid_x_min.item():.4f},"
+                f" {valid_x_max.item():.4f}], loss={x_band_loss.item():.6f}"
             )
             print(
-                f"    [On] Y: child_pos={child_pos[1].item():.4f}, valid_range=[{valid_y_min:.4f},"
-                f" {valid_y_max:.4f}], loss={y_band_loss.item():.6f}"
+                f"    [On] Y: child_pos={child_pos[1].item():.4f}, valid_range=[{valid_y_min.item():.4f},"
+                f" {valid_y_max.item():.4f}], loss={y_band_loss.item():.6f}"
+            )
+            print(
+                f"    [On] Z: child_pos={child_pos[2].item():.4f}, target={target_z.item():.4f},"
+                f" loss={z_loss.item():.6f}"
             )
-            print(f"    [On] Z: child_pos={child_pos[2].item():.4f}, target={target_z:.4f}, loss={z_loss.item():.6f}")
 
         total_loss = x_band_loss + y_band_loss + z_loss
         return relation.relation_loss_weight * total_loss
@@ -346,16 +349,16 @@ def compute_loss(
         """
         # Parent world extents from the world bounding box, expanded by clearance_m
         c = relation.clearance_m
-        parent_x_min = parent_world_bbox.min_point[0] - c
-        parent_x_max = parent_world_bbox.max_point[0] + c
-        parent_y_min = parent_world_bbox.min_point[1] - c
-        parent_y_max = parent_world_bbox.max_point[1] + c
-        parent_z_min = parent_world_bbox.min_point[2] - c
-        parent_z_max = parent_world_bbox.max_point[2] + c
+        parent_x_min = parent_world_bbox.min_point[0, 0] - c
+        parent_x_max = parent_world_bbox.max_point[0, 0] + c
+        parent_y_min = parent_world_bbox.min_point[0, 1] - c
+        parent_y_max = parent_world_bbox.max_point[0, 1] + c
+        parent_z_min = parent_world_bbox.min_point[0, 2] - c
+        parent_z_max = parent_world_bbox.max_point[0, 2] + c
 
         # Child world extents
-        child_world_min = child_pos + torch.tensor(child_bbox.min_point, dtype=child_pos.dtype, device=child_pos.device)
-        child_world_max = child_pos + torch.tensor(child_bbox.max_point, dtype=child_pos.dtype, device=child_pos.device)
+        child_world_min = child_pos + child_bbox.min_point[0]
+        child_world_max = child_pos + child_bbox.max_point[0]
 
         # 1. Per-axis overlap: zero when separated; else overlap length (default slope 1.0 gives length in m)
         overlap_x = interval_overlap_axis_loss(child_world_min[0], child_world_max[0], parent_x_min, parent_x_max)
@@ -369,15 +372,18 @@ def compute_loss(
         if self.debug:
             print(
                 f"    [NoCollision] X: overlap={overlap_x.item():.6f} (child_x=[{child_world_min[0].item():.4f},"
-                f" {child_world_max[0].item():.4f}], parent_x=[{parent_x_min:.4f}, {parent_x_max:.4f}])"
+                f" {child_world_max[0].item():.4f}], parent_x=[{parent_x_min.item():.4f},"
+                f" {parent_x_max.item():.4f}])"
             )
             print(
                 f"    [NoCollision] Y: overlap={overlap_y.item():.6f} (child_y=[{child_world_min[1].item():.4f},"
-                f" {child_world_max[1].item():.4f}], parent_y=[{parent_y_min:.4f}, {parent_y_max:.4f}])"
+                f" {child_world_max[1].item():.4f}], parent_y=[{parent_y_min.item():.4f},"
+                f" {parent_y_max.item():.4f}])"
             )
             print(
                 f"    [NoCollision] Z: overlap={overlap_z.item():.6f} (child_z=[{child_world_min[2].item():.4f},"
-                f" {child_world_max[2].item():.4f}], parent_z=[{parent_z_min:.4f}, {parent_z_max:.4f}])"
+                f" {child_world_max[2].item():.4f}], parent_z=[{parent_z_min.item():.4f},"
+                f" {parent_z_max.item():.4f}])"
             )
             print(f"    [NoCollision] volume={overlap_volume.item():.6f}, loss={total_loss.item():.6f}")
 

isaaclab_arena/relations/relation_solver.py

@@ -238,21 +238,34 @@ def _print_relation_debug(
 
     print(f"\n=== {obj.name} -> {type(relation).__name__}({relation.parent.name}) ===")
     print(f"  Child pos: ({child_pos[0].item():.4f}, {child_pos[1].item():.4f}, {child_pos[2].item():.4f})")
-    print(f"  Child bbox: min={child_bbox.min_point}, max={child_bbox.max_point}, size={child_bbox.size}")
+    print(f"  Child bbox: min={child_bbox.min_point[0].tolist()}, max={child_bbox.max_point[0].tolist()},"
+          f" size={child_bbox.size[0].tolist()}")
     print(f"  Parent pos: ({parent_pos[0].item():.4f}, {parent_pos[1].item():.4f}, {parent_pos[2].item():.4f})")
     print(
-        f"  Parent world bbox: min={parent_world_bbox.min_point}, max={parent_world_bbox.max_point},"
-        f" size={parent_world_bbox.size}"
+        f"  Parent world bbox: min={parent_world_bbox.min_point[0].tolist()},"
+        f" max={parent_world_bbox.max_point[0].tolist()}, size={parent_world_bbox.size[0].tolist()}"
     )
 
     # Child world extents
-    child_x_range = (child_pos[0].item() + child_bbox.min_point[0], child_pos[0].item() + child_bbox.max_point[0])
-    child_y_range = (child_pos[1].item() + child_bbox.min_point[1], child_pos[1].item() + child_bbox.max_point[1])
+    child_x_range = (
+        child_pos[0].item() + child_bbox.min_point[0, 0].item(),
+        child_pos[0].item() + child_bbox.max_point[0, 0].item(),
+    )
+    child_y_range = (
+        child_pos[1].item() + child_bbox.min_point[0, 1].item(),
+        child_pos[1].item() + child_bbox.max_point[0, 1].item(),
+    )
 
     print(f"  Child world X: [{child_x_range[0]:.4f}, {child_x_range[1]:.4f}]")
     print(f"  Child world Y: [{child_y_range[0]:.4f}, {child_y_range[1]:.4f}]")
-    print(f"  Parent world X: [{parent_world_bbox.min_point[0]:.4f}, {parent_world_bbox.max_point[0]:.4f}]")
-    print(f"  Parent world Y: [{parent_world_bbox.min_point[1]:.4f}, {parent_world_bbox.max_point[1]:.4f}]")
+    print(
+        f"  Parent world X: [{parent_world_bbox.min_point[0, 0].item():.4f},"
+        f" {parent_world_bbox.max_point[0, 0].item():.4f}]"
+    )
+    print(
+        f"  Parent world Y: [{parent_world_bbox.min_point[0, 1].item():.4f},"
+        f" {parent_world_bbox.max_point[0, 1].item():.4f}]"
+    )
     print(f"  Loss: {loss.item():.6f}")
 
 
@@ -270,5 +283,6 @@ def _print_unary_relation_debug(
     )
     print(f"\n=== {obj.name} -> {type(relation).__name__}({target_str}) ===")
     print(f"  Child pos: ({child_pos[0].item():.4f}, {child_pos[1].item():.4f}, {child_pos[2].item():.4f})")
-    print(f"  Child bbox: min={child_bbox.min_point}, max={child_bbox.max_point}, size={child_bbox.size}")
+    print(f"  Child bbox: min={child_bbox.min_point[0].tolist()}, max={child_bbox.max_point[0].tolist()},"
+          f" size={child_bbox.size[0].tolist()}")
     print(f"  Loss: {loss.item():.6f}")