Batch Support for Relation Solver

isaaclab_arena/assets/object_base.py

@@ -15,8 +15,9 @@
 
 from isaaclab_arena.assets.asset import Asset
 from isaaclab_arena.relations.relations import AtPosition, Relation, RelationBase
-from isaaclab_arena.terms.events import set_object_pose
-from isaaclab_arena.utils.pose import Pose, PoseRange
+from isaaclab_arena.terms.events import set_object_pose, set_object_pose_per_env
+from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+from isaaclab_arena.utils.pose import Pose, PosePerEnv, PoseRange
 from isaaclab_arena.utils.velocity import Velocity
 
 
@@ -42,38 +43,52 @@ def __init__(
             prim_path = "{ENV_REGEX_NS}/" + self.name
         self.prim_path = prim_path
         self.object_type = object_type
-        self.initial_pose: Pose | PoseRange | None = None
+        self.initial_pose: Pose | PoseRange | PosePerEnv | None = None
         self.initial_velocity: Velocity | None = None
         self.object_cfg = None
         self.event_cfg = None
         self.relations: list[RelationBase] = []
 
-    def get_initial_pose(self) -> Pose | PoseRange | None:
+    def get_initial_pose(self) -> Pose | PoseRange | PosePerEnv | None:
         """Return the current initial pose of this object.
 
         Subclasses may override to derive the pose from other sources
         (e.g. a parent asset), falling back to ``self.initial_pose``.
         """
         return self.initial_pose
 
+    @abstractmethod
+    def get_bounding_box(self) -> AxisAlignedBoundingBox:
+        """Get local bounding box (relative to object origin)."""
+        ...
+
+    @abstractmethod
+    def get_world_bounding_box(self) -> AxisAlignedBoundingBox:
+        """Get bounding box in world coordinates (local bbox rotated and translated)."""
+        ...
+
     def _get_initial_pose_as_pose(self) -> Pose | None:
         """Return a single ``Pose`` suitable for *init_state* and bounding-box calculations.
 
         If the initial pose is a ``PoseRange``, its midpoint is returned.
+        If the initial pose is a ``PosePerEnv``, the first environment's pose is returned.
         If the initial pose is ``None``, ``None`` is returned.
         """
         initial_pose = self.get_initial_pose()
         if initial_pose is None:
             return None
+        if isinstance(initial_pose, PosePerEnv):
+            return initial_pose.poses[0]
         if isinstance(initial_pose, PoseRange):
             return initial_pose.get_midpoint()
         return initial_pose
 
-    def set_initial_pose(self, pose: Pose | PoseRange) -> None:
+    def set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None:
         """Set / override the initial pose and rebuild derived configs.
 
         Args:
-            pose: A fixed ``Pose`` or a ``PoseRange`` (randomised on reset).
+            pose: A fixed ``Pose``, a ``PoseRange`` (randomised on reset),
+                or a ``PosePerEnv`` (distinct pose per environment).
         """
         self.initial_pose = pose
         initial_pose = self._get_initial_pose_as_pose()
@@ -116,7 +131,16 @@ def _init_event_cfg(self) -> EventTermCfg | None:
             return None
 
         initial_pose = self.get_initial_pose()
-        if isinstance(initial_pose, PoseRange):
+        if isinstance(initial_pose, PosePerEnv):
+            return EventTermCfg(
+                func=set_object_pose_per_env,
+                mode="reset",
+                params={
+                    "asset_cfg": SceneEntityCfg(self.name),
+                    "pose_list": initial_pose.poses,
+                },
+            )
+        elif isinstance(initial_pose, PoseRange):
             return EventTermCfg(
                 func=randomize_object_pose,
                 mode="reset",

isaaclab_arena/environments/arena_env_builder.py

@@ -27,6 +27,7 @@
 from isaaclab_arena.metrics.recorder_manager_utils import metrics_to_recorder_manager_cfg
 from isaaclab_arena.relations.object_placer import ObjectPlacer
 from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
+from isaaclab_arena.relations.placement_result import MultiEnvPlacementResult
 from isaaclab_arena.relations.relations import IsAnchor, NoCollision
 from isaaclab_arena.tasks.no_task import NoTask
 from isaaclab_arena.utils.configclass import combine_configclass_instances
@@ -99,11 +100,24 @@ def _solve_relations(self) -> None:
         self._add_pairwise_no_collision(objects_with_relations)
 
         # Run the ObjectPlacer (default on_relation_z_tolerance_m accommodates solver residual).
+        # Positions are applied to objects via set_initial_pose (single-env: Pose/PoseRange,
+        # multi-env: PosePerEnv), so each object's event_cfg handles its own reset.
         placement_seed = getattr(self.args, "placement_seed", None)
         placer = ObjectPlacer(params=ObjectPlacerParams(placement_seed=placement_seed))
-        result = placer.place(objects=objects_with_relations)
-
-        if result.success:
+        num_envs = self.args.num_envs
+        result = placer.place(objects_with_relations, num_envs=num_envs)
+
+        # Log outcome
+        if isinstance(result, MultiEnvPlacementResult):
+            n_succeeded = sum(1 for r in result.results if r.success)
+            if n_succeeded == num_envs:
+                print(f"Relation solving succeeded for all {num_envs} env(s) after {result.attempts} attempt(s)")
+            else:
+                print(
+                    f"Relation solving: {n_succeeded}/{num_envs} env(s) passed validation after"
+                    f" {result.attempts} attempt(s)."
+                )
+        elif result.success:
             print(f"Relation solving succeeded after {result.attempts} attempt(s)")
         else:
             print(f"Relation solving not completed after {result.attempts} attempt(s)")