supoort for rotation, verification of algorithm

Signed-off-by: zhx06 <zihaox@nvidia.com>

Author: zhx06

isaaclab_arena/environments/relation_solver_interface.py

@@ -52,6 +52,7 @@ def solve_and_apply_relation_placement(
         print("No objects with relations found in scene. Skipping relation solving.")
         return None
 
+    assert collision_mode in ("bbox", "mesh"), f"Invalid collision_mode '{collision_mode}', expected 'bbox' or 'mesh'"
     mode = CollisionMode.MESH if collision_mode == "mesh" else CollisionMode.BBOX
     placer_params = ObjectPlacerParams(
         placement_seed=placement_seed,

isaaclab_arena/relations/collision_mode.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from enum import Enum
+
+
+class CollisionMode(Enum):
+    """Selects which collision detection method the solver uses for no-overlap constraints."""
+
+    BBOX = "bbox"
+    """Axis-aligned bounding box overlap volume (fast, conservative)."""
+
+    MESH = "mesh"
+    """Sphere-to-SDF queries against actual mesh geometry (accurate, slower)."""

isaaclab_arena/relations/object_placer.py

@@ -11,10 +11,10 @@
 from typing import TYPE_CHECKING
 
 from isaaclab_arena.relations.bounding_box_helpers import assign_variants_for_envs, build_per_env_bounding_boxes
+from isaaclab_arena.relations.collision_mode import CollisionMode
 from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
 from isaaclab_arena.relations.placement_result import MultiEnvPlacementResult, PlacementResult
 from isaaclab_arena.relations.relation_solver import RelationSolver
-from isaaclab_arena.relations.relation_solver_params import CollisionMode
 from isaaclab_arena.relations.relations import (
     IsAnchor,
     On,
@@ -23,7 +23,7 @@
     get_anchor_objects,
 )
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-from isaaclab_arena.utils.pose import Pose, PosePerEnv, rotate_quat_by_yaw, wrap_angle_to_pi
+from isaaclab_arena.utils.pose import Pose, PosePerEnv, rotate_quat_by_yaw, wrap_angle_to_pi, yaw_from_quat_xyzw
 from isaaclab_arena.utils.random import get_random_rotation
 
 if TYPE_CHECKING:
@@ -673,9 +673,35 @@ def _validate_no_overlap_mesh(
                         if mesh is None and obj.name not in warned_no_mesh:
                             warned_no_mesh.add(obj.name)
                             print(
-                                f"  [NoCollision] MESH mode: '{obj.name}' has no collision mesh, skipping mesh"
-                                " validation"
+                                f"  [NoCollision] MESH mode: '{obj.name}' has no collision mesh,"
+                                " falling back to AABB validation for this pair"
                             )
+                    # Fall back to AABB overlap check (matching the solver AABB fallback).
+                    for obj in (a, b):
+                        if id(obj) in anchor_ids:
+                            pose = obj.get_initial_pose()
+                            if pose is not None and hasattr(pose, "rotation_xyzw"):
+                                qx, qy = pose.rotation_xyzw[0], pose.rotation_xyzw[1]
+                                assert abs(qx) < 1e-6 and abs(qy) < 1e-6, (
+                                    f"AABB fallback requires anchor '{obj.name}' to have pure-Z rotation, "
+                                    f"got rotation_xyzw={pose.rotation_xyzw}"
+                                )
+                    a_pos = torch.tensor(positions[a], dtype=torch.float32)
+                    b_pos = torch.tensor(positions[b], dtype=torch.float32)
+                    a_bbox = a.get_bounding_box()
+                    b_bbox = b.get_bounding_box()
+                    a_yaw = ObjectPlacer._effective_yaw(a, orientations)
+                    b_yaw = ObjectPlacer._effective_yaw(b, orientations)
+                    if a_yaw != 0.0:
+                        a_bbox = a_bbox.rotated_around_z(a_yaw)
+                    if b_yaw != 0.0:
+                        b_bbox = b_bbox.rotated_around_z(b_yaw)
+                    a_world = a_bbox.translated(a_pos)
+                    b_world = b_bbox.translated(b_pos)
+                    if a_world.overlaps(b_world, margin=tolerance).item():
+                        if self.params.verbose:
+                            print(f"  AABB overlap between '{a.name}' and '{b.name}' (mesh unavailable)")
+                        return False
                     continue
 
                 a_pos = torch.tensor(positions[a], dtype=torch.float32)
@@ -718,10 +744,7 @@ def _effective_yaw(obj: ObjectBase, orientations: dict[ObjectBase, float] | None
         pose = obj.get_initial_pose()
         if pose is None or not hasattr(pose, "rotation_xyzw"):
             return 0.0
-        qx, qy, qz, qw = pose.rotation_xyzw
-        if abs(qx) > 1e-6 or abs(qy) > 1e-6:
-            return 0.0
-        return 2.0 * math.atan2(qz, qw)
+        return yaw_from_quat_xyzw(pose.rotation_xyzw)
 
     @staticmethod
     def _centers_in_target_frame(

isaaclab_arena/relations/relation_loss_strategies.py

@@ -5,12 +5,14 @@
 
 from __future__ import annotations
 
+import math
 import torch
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import IntEnum
 from typing import TYPE_CHECKING
 
+from isaaclab_arena.relations.collision_mode import CollisionMode
 from isaaclab_arena.relations.loss_primitives import (
     interval_overlap_axis_loss,
     linear_band_loss,
@@ -20,7 +22,11 @@
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
+    import trimesh
+
+    from isaaclab_arena.assets.object_base import ObjectBase
     from isaaclab_arena.relations.relations import AtPosition, NextTo, NotNextTo, On, PositionLimits, Relation
+    from isaaclab_arena.relations.warp_mesh_manager import WarpMeshManager
 
 from isaaclab_arena.relations.relations import Side
 
@@ -427,7 +433,7 @@ def __init__(
         self,
         slope: float = 10.0,
         debug: bool = False,
-        collision_mode=None,
+        collision_mode: CollisionMode | None = None,
         num_spheres: int = 30,
     ):
         """
@@ -437,15 +443,12 @@ def __init__(
             collision_mode: CollisionMode enum value (BBOX or MESH). Defaults to BBOX.
             num_spheres: Number of spheres for mesh decomposition (MESH mode only).
         """
-        from isaaclab_arena.relations.relation_solver_params import CollisionMode
-
         self.slope = slope
         self.debug = debug
         self._mode = collision_mode if collision_mode is not None else CollisionMode.BBOX
         self._num_spheres = num_spheres
-        self._CollisionMode = CollisionMode
         self._warned_no_mesh: set[str] = set()
-        self._mesh_managers: dict[str, object] = {}  # keyed by device string
+        self._mesh_managers: dict[str, WarpMeshManager] = {}
 
         if self._mode == CollisionMode.MESH:
             try:
@@ -461,9 +464,11 @@ def compute_loss(
         child_pos: torch.Tensor,
         child_bbox: AxisAlignedBoundingBox,
         parent_world_bbox: AxisAlignedBoundingBox,
-        child_obj=None,
-        parent_obj=None,
+        child_obj: ObjectBase | None = None,
+        parent_obj: ObjectBase | None = None,
         parent_pos: torch.Tensor | None = None,
+        child_yaw: float = 0.0,
+        parent_yaw: float = 0.0,
     ) -> torch.Tensor:
         """Compute collision loss, dispatching based on mode and mesh availability.
 
@@ -475,16 +480,26 @@ def compute_loss(
             child_obj: Object with get_collision_mesh() (MESH mode).
             parent_obj: Object with get_collision_mesh() (MESH mode).
             parent_pos: Parent position tensor, or None for anchors (MESH mode).
+            child_yaw: Z-yaw (radians) of the child object (MESH mode).
+            parent_yaw: Z-yaw (radians) of the parent object (MESH mode).
 
         Returns:
             Loss tensor of shape (N,).
         """
-        if self._mode == self._CollisionMode.MESH and child_obj is not None and parent_obj is not None:
+        if self._mode == CollisionMode.MESH and child_obj is not None and parent_obj is not None:
             child_mesh = child_obj.get_collision_mesh()
             parent_mesh = parent_obj.get_collision_mesh()
             if child_mesh is not None and parent_mesh is not None:
                 return self._compute_mesh_loss(
-                    clearance_m, child_pos, child_obj, child_mesh, parent_pos, parent_obj, parent_mesh
+                    clearance_m,
+                    child_pos,
+                    child_obj,
+                    child_mesh,
+                    parent_pos,
+                    parent_obj,
+                    parent_mesh,
+                    child_yaw=child_yaw,
+                    parent_yaw=parent_yaw,
                 )
             for obj, mesh in [(child_obj, child_mesh), (parent_obj, parent_mesh)]:
                 name = getattr(obj, "name", "?")
@@ -544,7 +559,7 @@ def _compute_aabb_loss(
 
         return total_loss.squeeze(0) if single_input else total_loss
 
-    def _get_mesh_manager(self, device: str = "cuda:0"):
+    def _get_mesh_manager(self, device: str = "cuda:0") -> WarpMeshManager:
         """Return a cached WarpMeshManager for the given device."""
         if device not in self._mesh_managers:
             from isaaclab_arena.relations.warp_mesh_manager import WarpMeshManager
@@ -556,13 +571,15 @@ def _compute_mesh_loss(
         self,
         clearance_m: float,
         child_pos: torch.Tensor,
-        child_obj,
-        child_mesh,
+        child_obj: ObjectBase,
+        child_mesh: trimesh.Trimesh,
         parent_pos: torch.Tensor | None,
-        parent_obj,
-        parent_mesh,
+        parent_obj: ObjectBase,
+        parent_mesh: trimesh.Trimesh,
+        child_yaw: float = 0.0,
+        parent_yaw: float = 0.0,
     ) -> torch.Tensor:
-        """Sphere-to-SDF penetration loss using mesh geometry."""
+        """Per-pair sphere-to-SDF penetration loss."""
         from isaaclab_arena.relations.warp_sdf_kernels import sphere_penetration_loss
 
         single_input = child_pos.dim() == 1
@@ -577,12 +594,11 @@ def _compute_mesh_loss(
             assert isinstance(
                 pose, Pose
             ), f"Anchor '{getattr(parent_obj, 'name', '?')}' must have a fixed Pose for mesh collision"
-            identity = (0.0, 0.0, 0.0, 1.0)
-            assert pose.rotation_xyzw == identity, (
-                f"Mesh collision with rotated anchor '{getattr(parent_obj, 'name', '?')}' "
-                f"is not yet supported (rotation={pose.rotation_xyzw})"
-            )
             parent_pos = torch.tensor(pose.position_xyz, dtype=child_pos.dtype, device=child_pos.device)
+            if parent_yaw == 0.0:
+                from isaaclab_arena.utils.pose import yaw_from_quat_xyzw
+
+                parent_yaw = yaw_from_quat_xyzw(pose.rotation_xyzw)
 
         assert parent_pos is not None
         parent_pos_resolved: torch.Tensor = parent_pos
@@ -596,14 +612,31 @@ def _compute_mesh_loss(
         radii = spheres[:, 3]
         warp_mesh = manager.get_warp_mesh(parent_mesh, obj=parent_obj)
 
+        # Rotate child sphere centers by net_yaw = child_yaw - parent_yaw.
+        net_yaw = child_yaw - parent_yaw
+        if net_yaw != 0.0:
+            cos_n = math.cos(net_yaw)
+            sin_n = math.sin(net_yaw)
+            rx = centers_local[:, 0] * cos_n - centers_local[:, 1] * sin_n
+            ry = centers_local[:, 0] * sin_n + centers_local[:, 1] * cos_n
+            centers_local = torch.stack([rx, ry, centers_local[:, 2]], dim=-1)
+
         batch_size = child_pos.shape[0]
         parent_pos_resolved = parent_pos_resolved.expand(batch_size, -1)
         total_loss = torch.zeros(batch_size, device=device, dtype=child_pos.dtype)
 
         for b in range(batch_size):
-            centers_world = centers_local + child_pos[b] - parent_pos_resolved[b]
+            offset = child_pos[b] - parent_pos_resolved[b]
+            # Rotate offset into the parent's local frame.
+            if parent_yaw != 0.0:
+                cos_p = math.cos(-parent_yaw)
+                sin_p = math.sin(-parent_yaw)
+                ox = offset[0] * cos_p - offset[1] * sin_p
+                oy = offset[0] * sin_p + offset[1] * cos_p
+                offset = torch.stack([ox, oy, offset[2]])
+            centers_in_parent = centers_local + offset
             total_loss[b] = self.slope * sphere_penetration_loss(
-                centers_world, radii, warp_mesh, clearance_m=clearance_m
+                centers_in_parent, radii, warp_mesh, clearance_m=clearance_m
             )
 
         return total_loss.squeeze(0) if single_input else total_loss

isaaclab_arena/relations/relation_solver.py

@@ -5,16 +5,16 @@
 
 from __future__ import annotations
 
-import math
 import torch
 from typing import TYPE_CHECKING, cast
 
+from isaaclab_arena.relations.collision_mode import CollisionMode
 from isaaclab_arena.relations.relation_loss_strategies import (
     NoCollisionLossStrategy,
     RelationLossStrategy,
     UnaryRelationLossStrategy,
 )
-from isaaclab_arena.relations.relation_solver_params import CollisionMode, RelationSolverParams
+from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
 from isaaclab_arena.relations.relation_solver_state import RelationSolverState
 from isaaclab_arena.relations.relations import On, Relation, RelationBase, UnaryRelation
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
@@ -182,25 +182,29 @@ def _compute_no_overlap_loss(
             Per-environment loss tensor of shape (batch_size,).
         """
         if self.params.collision_mode == CollisionMode.MESH:
-            return self._compute_no_overlap_loss_mesh(state, debug)
+            mesh_loss = self._compute_no_overlap_loss_mesh(state, debug)
+            aabb_loss = self._compute_no_overlap_loss_aabb(state, debug, skip_mesh_pairs=True)
+            return mesh_loss + aabb_loss
         else:
             return self._compute_no_overlap_loss_aabb(state, debug)
 
     def _compute_no_overlap_loss_aabb(
         self,
         state: RelationSolverState,
         debug: bool,
+        skip_mesh_pairs: bool = False,
     ) -> torch.Tensor:
-        """Per-pair AABB collision loss (used when collision_mode != MESH)."""
+        """Per-pair AABB collision loss.
+
+        When skip_mesh_pairs=True (used as AABB fallback in MESH mode), only
+        processes pairs where at least one object lacks a collision mesh.
+        """
         device = state.device
         total_loss = torch.zeros(state.batch_size, device=device, dtype=torch.float32)
 
         non_anchor_objects = state.optimizable_objects
         anchor_objects = list(state.anchor_objects)
 
-        # Skip no-overlap for On-linked pairs: the On constraint already forces
-        # vertical contact, and penalizing overlap here would fight it, causing
-        # oscillation between "push apart" and "keep on surface" gradients.
         on_pairs: set[tuple[int, int]] = set()
         for obj in [*non_anchor_objects, *anchor_objects]:
             for rel in obj.get_relations():
@@ -215,6 +219,12 @@ def _compute_no_overlap_loss_aabb(
             for anchor in anchor_objects:
                 if (id(child), id(anchor)) in on_pairs:
                     continue
+                if (
+                    skip_mesh_pairs
+                    and child.get_collision_mesh() is not None
+                    and anchor.get_collision_mesh() is not None
+                ):
+                    continue
                 anchor_world_bbox = anchor.get_world_bounding_box().to(device)
                 loss = self._no_collision_strategy.compute_loss(
                     clearance_m=self.params.clearance_m,
@@ -233,6 +243,12 @@ def _compute_no_overlap_loss_aabb(
                 other = non_anchor_objects[j]
                 if (id(child), id(other)) in on_pairs:
                     continue
+                if (
+                    skip_mesh_pairs
+                    and child.get_collision_mesh() is not None
+                    and other.get_collision_mesh() is not None
+                ):
+                    continue
                 other_pos = state.get_position(other)
                 other_bbox = state.get_bbox(other)
 
@@ -302,7 +318,7 @@ def _build_vectorized_cache(
 
         import warp as wp
 
-        from isaaclab_arena.utils.pose import Pose
+        from isaaclab_arena.utils.pose import Pose, yaw_from_quat_xyzw
 
         centers_list: list[torch.Tensor] = []
         radii_list: list[torch.Tensor] = []
@@ -352,14 +368,13 @@ def _build_vectorized_cache(
                     p_bbox_max = parent_bbox.max_point.to(device)
                     pose = anchor.get_initial_pose()
                     assert pose is not None and isinstance(pose, Pose)
-                    qx, qy, qz, qw = pose.rotation_xyzw
-                    assert abs(qx) < 1e-6 and abs(qy) < 1e-6, (
+                    assert abs(pose.rotation_xyzw[0]) < 1e-6 and abs(pose.rotation_xyzw[1]) < 1e-6, (
                         f"MESH collision requires anchor '{anchor.name}' to have identity or "
                         f"pure-Z rotation, got rotation_xyzw={pose.rotation_xyzw}. "
                         "Roll/pitch anchors are not supported in MESH mode."
                     )
                     anchor_pos = torch.tensor(pose.position_xyz, dtype=torch.float32, device=device)
-                    anchor_yaw = 2.0 * math.atan2(qz, qw)
+                    anchor_yaw = yaw_from_quat_xyzw(pose.rotation_xyzw)
 
                     n_spheres = child_centers_local.shape[0]
                     mesh_key = id(warp_mesh)