From 7f826b29d0dd059c2b41df4d9a169a4a87b7019c Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 13 Apr 2026 14:34:02 -0700
Subject: [PATCH 01/46] support for heterogeneous objects

---
 isaaclab_arena/assets/dummy_object.py         |   8 +
 isaaclab_arena/assets/object_base.py          |  20 ++
 isaaclab_arena/assets/object_set.py           |  46 +++
 isaaclab_arena/relations/object_placer.py     | 175 +++++++++---
 isaaclab_arena/relations/relation_solver.py   |  43 ++-
 .../tests/test_heterogeneous_placement.py     | 266 ++++++++++++++++++
 6 files changed, 519 insertions(+), 39 deletions(-)
 create mode 100644 isaaclab_arena/tests/test_heterogeneous_placement.py

diff --git a/isaaclab_arena/assets/dummy_object.py b/isaaclab_arena/assets/dummy_object.py
index 028687e63..6551ecd6f 100644
--- a/isaaclab_arena/assets/dummy_object.py
+++ b/isaaclab_arena/assets/dummy_object.py
@@ -42,6 +42,14 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get local bounding box (relative to object origin)."""
         return self.bounding_box
 
+    def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
+        """Get per-environment local bounding boxes (expanded from single bbox)."""
+        bbox = self.get_bounding_box()
+        return AxisAlignedBoundingBox(
+            min_point=bbox.min_point.expand(num_envs, 3),
+            max_point=bbox.max_point.expand(num_envs, 3),
+        )
+
     def get_world_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get bounding box in world coordinates (local bbox rotated and translated).
 
diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py
index 075d8f1c9..35ecc58b1 100644
--- a/isaaclab_arena/assets/object_base.py
+++ b/isaaclab_arena/assets/object_base.py
@@ -73,6 +73,26 @@ def get_world_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get bounding box in world coordinates (local bbox rotated and translated)."""
         ...
 
+    def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
+        """Get per-environment local bounding boxes.
+
+        For homogeneous objects the single local bbox is expanded to ``(num_envs, 3)``.
+        ``RigidObjectSet`` overrides this to return the actual bbox of each env's
+        variant, enabling heterogeneous placement.
+
+        Args:
+            num_envs: Number of environments.
+
+        Returns:
+            ``AxisAlignedBoundingBox`` with ``min_point`` / ``max_point`` of shape
+            ``(num_envs, 3)``.
+        """
+        bbox = self.get_bounding_box()
+        return AxisAlignedBoundingBox(
+            min_point=bbox.min_point.expand(num_envs, 3),
+            max_point=bbox.max_point.expand(num_envs, 3),
+        )
+
     def _get_initial_pose_as_pose(self) -> Pose | None:
         """Return a single ``Pose`` suitable for *init_state* and bounding-box calculations.
 
diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 14d6dea12..6df45acf6 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -3,6 +3,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import torch
+
 import isaaclab.sim as sim_utils
 from isaaclab.assets import RigidObjectCfg
 from isaaclab.sensors.contact_sensor.contact_sensor_cfg import ContactSensorCfg
@@ -67,6 +69,7 @@ def __init__(
 
         self.objects: list[Object] = objects
         self.random_choice = random_choice
+        self.heterogeneous_bbox: bool = len(objects) > 1
 
         # Set default prim_path if not provided
         if prim_path is None:
@@ -90,6 +93,49 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """
         return max(self.objects, key=lambda obj: obj.get_bounding_box().size[0, 2].item()).get_bounding_box()
 
+    def get_variant_indices(self, num_envs: int) -> list[int]:
+        """Return which member object index is assigned to each environment.
+
+        When ``random_choice`` is False the mapping is round-robin
+        (``env_idx % len(objects)``).  When True, a random permutation is
+        generated (and cached so repeated calls with the same ``num_envs``
+        are deterministic within a session).
+
+        Args:
+            num_envs: Number of environments.
+
+        Returns:
+            List of length ``num_envs`` with indices into ``self.objects``.
+        """
+        n = len(self.objects)
+        if not self.random_choice:
+            return [i % n for i in range(num_envs)]
+
+        if not hasattr(self, "_cached_variant_indices") or len(self._cached_variant_indices) != num_envs:
+            self._cached_variant_indices = [int(i) for i in torch.randint(n, (num_envs,)).tolist()]
+        return self._cached_variant_indices
+
+    def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
+        """Get the actual bounding box for each env's variant.
+
+        Unlike ``get_bounding_box()`` (which uses a max-z heuristic), this
+        returns the real local bbox of the variant assigned to each env,
+        enabling correct collision-free placement for heterogeneous scenes.
+
+        Args:
+            num_envs: Number of environments.
+
+        Returns:
+            ``AxisAlignedBoundingBox`` with ``min_point`` / ``max_point`` of
+            shape ``(num_envs, 3)``.
+        """
+        variant_indices = self.get_variant_indices(num_envs)
+        member_bboxes = [obj.get_bounding_box() for obj in self.objects]
+
+        min_pts = torch.stack([member_bboxes[idx].min_point[0] for idx in variant_indices], dim=0)
+        max_pts = torch.stack([member_bboxes[idx].max_point[0] for idx in variant_indices], dim=0)
+        return AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts)
+
     def get_contact_sensor_cfg(self, contact_against_object: ObjectBase | None = None) -> ContactSensorCfg:
         # We assume that by here, our USDs have been modified to be compatible with each other
         # and we can use the first USD path to find the shallowest rigid body.
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index f11a1030a..7eaa7d82f 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -27,6 +27,16 @@
     from isaaclab_arena.assets.object_base import ObjectBase
 
 
+def _is_heterogeneous(obj: ObjectBase) -> bool:
+    """Return True if *obj* provides per-env variant geometry.
+
+    ``RigidObjectSet`` (and test doubles) set ``heterogeneous_bbox = True``
+    to signal that ``get_bounding_box_per_env`` returns different bboxes
+    across environments.
+    """
+    return getattr(obj, "heterogeneous_bbox", False)
+
+
 @dataclass
 class PlacementCandidate:
     """A single solver result, ranked and selected in ObjectPlacer.place()."""
@@ -119,12 +129,44 @@ def place(
         if self.params.placement_seed is not None:
             generator = torch.Generator()
 
-        # Pool-based placement: generate all candidates in one batched call,
-        # then pick the best num_results (environments are homogeneous so any
-        # valid solution can serve any environment).
         num_results = num_envs if result_per_env else 1
-        num_candidates = self.params.max_placement_attempts * num_results
+        max_attempts = self.params.max_placement_attempts
+        num_candidates = max_attempts * num_results
+
+        # Detect heterogeneous objects (e.g. RigidObjectSet with per-env variants).
+        heterogeneous = result_per_env and any(_is_heterogeneous(obj) for obj in objects)
+
+        if heterogeneous:
+            results_per_env = self._place_heterogeneous(
+                objects, anchor_objects_set, num_envs, max_attempts, num_candidates, generator
+            )
+        else:
+            results_per_env = self._place_homogeneous(
+                objects, anchor_objects_set, num_results, max_attempts, num_candidates, generator
+            )
+
+        final_per_env = [r.positions for r in results_per_env]
+        if self.params.apply_positions_to_objects:
+            self._apply_positions(final_per_env, anchor_objects_set)
+
+        if num_results == 1:
+            return results_per_env[0]
+        return MultiEnvPlacementResult(results=results_per_env)
 
+    # ------------------------------------------------------------------
+    # Placement strategies
+    # ------------------------------------------------------------------
+
+    def _place_homogeneous(
+        self,
+        objects: list[ObjectBase],
+        anchor_objects_set: set[ObjectBase],
+        num_results: int,
+        max_attempts: int,
+        num_candidates: int,
+        generator: torch.Generator | None,
+    ) -> list[PlacementResult]:
+        """Pool-based placement: any valid solution can serve any environment."""
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
         for candidate_idx in range(num_candidates):
             if generator is not None:
@@ -135,51 +177,110 @@ def place(
         assert self._solver.last_loss_per_env is not None
         all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
 
-        all_candidates: list[PlacementCandidate] = []
-        for idx in range(num_candidates):
-            loss = all_losses[idx]
-            is_valid = self._validate_placement(all_positions[idx])
-            all_candidates.append(PlacementCandidate(loss, all_positions[idx], is_valid))
-
-        # Sort: valid solutions first (by loss), then invalid (by loss)
-        all_candidates.sort(key=lambda candidate: (not candidate.is_valid, candidate.loss))
+        all_candidates = [
+            PlacementCandidate(all_losses[i], all_positions[i], self._validate_placement(all_positions[i]))
+            for i in range(num_candidates)
+        ]
+        all_candidates.sort(key=lambda c: (not c.is_valid, c.loss))
         selected = all_candidates[:num_results]
 
-        n_valid = sum(1 for candidate in selected if candidate.is_valid)
         if self.params.verbose:
-            total_valid = sum(1 for candidate in all_candidates if candidate.is_valid)
-            finite_losses = [candidate.loss for candidate in all_candidates if math.isfinite(candidate.loss)]
+            total_valid = sum(1 for c in all_candidates if c.is_valid)
+            finite_losses = [c.loss for c in all_candidates if math.isfinite(c.loss)]
             mean_loss = sum(finite_losses) / len(finite_losses) if finite_losses else float("inf")
+            n_valid = sum(1 for c in selected if c.is_valid)
             print(
                 f"Solved {num_candidates} candidates in one batch: mean loss = {mean_loss:.6f},"
                 f" {total_valid} valid, selected best {num_results} ({n_valid} valid)"
             )
 
-        final_per_env: list[dict[ObjectBase, tuple[float, float, float]]] = [
-            candidate.positions for candidate in selected
+        return [
+            PlacementResult(success=c.is_valid, positions=c.positions, final_loss=c.loss, attempts=max_attempts)
+            for c in selected
         ]
-        results_per_env = [
-            PlacementResult(
-                success=candidate.is_valid,
-                positions=candidate.positions,
-                final_loss=candidate.loss,
-                attempts=self.params.max_placement_attempts,
+
+    def _place_heterogeneous(
+        self,
+        objects: list[ObjectBase],
+        anchor_objects_set: set[ObjectBase],
+        num_envs: int,
+        max_attempts: int,
+        num_candidates: int,
+        generator: torch.Generator | None,
+    ) -> list[PlacementResult]:
+        """Per-env placement: each candidate is tied to its env's object variants.
+
+        Batch layout: candidates [e * max_attempts : (e+1) * max_attempts] belong
+        to env *e*. Per-row bboxes reflect each env's actual variant geometry.
+        """
+        # Build per-env bboxes (num_envs, 3) for every object.
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {
+            obj: obj.get_bounding_box_per_env(num_envs) for obj in objects
+        }
+
+        # Expand into per-row bboxes (num_candidates, 3): repeat each env's
+        # bbox max_attempts times so rows [e*A:(e+1)*A] share env e's geometry.
+        bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] = {}
+        for obj, bbox in env_bboxes.items():
+            # bbox.min_point is (num_envs, 3) → repeat_interleave → (num_candidates, 3)
+            min_pt = bbox.min_point.repeat_interleave(max_attempts, dim=0)
+            max_pt = bbox.max_point.repeat_interleave(max_attempts, dim=0)
+            bboxes_per_row[obj] = AxisAlignedBoundingBox(min_point=min_pt, max_point=max_pt)
+
+        # Generate initial positions; each candidate uses its env's bbox.
+        initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
+        for candidate_idx in range(num_candidates):
+            env_idx = candidate_idx // max_attempts
+            if generator is not None:
+                generator.manual_seed(self.params.placement_seed + candidate_idx)
+            # Slice single-env bboxes for this candidate's env.
+            env_child_bboxes = {
+                obj: AxisAlignedBoundingBox(
+                    min_point=env_bboxes[obj].min_point[env_idx : env_idx + 1],
+                    max_point=env_bboxes[obj].max_point[env_idx : env_idx + 1],
+                )
+                for obj in objects
+            }
+            initial_positions.append(
+                self._generate_initial_positions(objects, anchor_objects_set, generator, child_bboxes=env_child_bboxes)
             )
-            for candidate in selected
-        ]
 
-        if self.params.apply_positions_to_objects:
-            self._apply_positions(final_per_env, anchor_objects_set)
+        all_positions = self._solver.solve(objects, initial_positions, bboxes_per_row=bboxes_per_row)
+        assert self._solver.last_loss_per_env is not None
+        all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
 
-        if num_results == 1:
-            return results_per_env[0]
-        return MultiEnvPlacementResult(results=results_per_env)
+        # Select best candidate per env.
+        results: list[PlacementResult] = []
+        for env_idx in range(num_envs):
+            start = env_idx * max_attempts
+            env_candidates = [
+                PlacementCandidate(
+                    all_losses[start + j],
+                    all_positions[start + j],
+                    self._validate_placement(all_positions[start + j]),
+                )
+                for j in range(max_attempts)
+            ]
+            env_candidates.sort(key=lambda c: (not c.is_valid, c.loss))
+            best = env_candidates[0]
+            results.append(
+                PlacementResult(
+                    success=best.is_valid, positions=best.positions, final_loss=best.loss, attempts=max_attempts
+                )
+            )
+
+        if self.params.verbose:
+            n_valid = sum(1 for r in results if r.success)
+            print(f"Heterogeneous placement: {n_valid}/{num_envs} env(s) valid")
+
+        return results
 
     def _generate_initial_positions(
         self,
         objects: list[ObjectBase],
         anchor_objects: set[ObjectBase],
         generator: torch.Generator | None = None,
+        child_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> dict[ObjectBase, tuple[float, float, float]]:
         """Generate initial positions for all objects.
 
@@ -190,6 +291,9 @@ def _generate_initial_positions(
         Args:
             generator: Optional RNG generator for reproducible sampling. When None,
                 uses PyTorch's global RNG.
+            child_bboxes: Optional per-object bbox overrides with shape ``(1, 3)``.
+                Used by heterogeneous placement to supply the correct variant
+                bbox when computing On-guided initial positions.
 
         Returns:
             Dictionary mapping all objects to their starting positions.
@@ -204,7 +308,10 @@ def _generate_initial_positions(
             if obj in anchor_objects:
                 positions[obj] = obj.get_initial_pose().position_xyz
             elif any(isinstance(r, On) for r in obj.get_relations()):
-                positions[obj] = self._compute_on_guided_position(obj, anchor_objects, anchor_bbox, generator)
+                bbox_override = child_bboxes.get(obj) if child_bboxes else None
+                positions[obj] = self._compute_on_guided_position(
+                    obj, anchor_objects, anchor_bbox, generator, child_bbox=bbox_override
+                )
             else:
                 positions[obj] = (cx, cy, cz)
         return positions
@@ -237,6 +344,7 @@ def _compute_on_guided_position(
         anchor_objects: set[ObjectBase],
         anchor_bbox: AxisAlignedBoundingBox,
         generator: torch.Generator | None = None,
+        child_bbox: AxisAlignedBoundingBox | None = None,
     ) -> tuple[float, float, float]:
         """Compute an initial position for an object with an On relation.
 
@@ -246,10 +354,13 @@ def _compute_on_guided_position(
         Args:
             generator: Optional RNG generator for reproducible sampling. When None,
                 uses PyTorch's global RNG.
+            child_bbox: Optional bbox override for the child object. When ``None``,
+                ``obj.get_bounding_box()`` is used.
         """
         on_relation = next(r for r in obj.get_relations() if isinstance(r, On))
         parent_bbox = self._get_on_parent_world_bbox(on_relation.parent, anchor_objects, anchor_bbox)
-        child_bbox = obj.get_bounding_box()
+        if child_bbox is None:
+            child_bbox = obj.get_bounding_box()
 
         x = self._sample_axis_position(
             parent_bbox.min_point[0, 0],
diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index 066dfa069..ebc042b14 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -15,7 +15,8 @@
 )
 from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
 from isaaclab_arena.relations.relation_solver_state import RelationSolverState
-from isaaclab_arena.relations.relations import Relation, RelationBase, UnaryRelation
+from isaaclab_arena.relations.relations import AtPosition, Relation, RelationBase, UnaryRelation
+from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
     from isaaclab_arena.assets.object_base import ObjectBase
@@ -66,12 +67,32 @@ def _get_strategy(self, relation: RelationBase) -> RelationLossStrategy | UnaryR
             )
         return strategy
 
-    def _compute_total_loss(self, state: RelationSolverState, debug: bool = False) -> torch.Tensor:
+    def _get_bbox(
+        self,
+        obj: ObjectBase,
+        device: torch.device | None,
+        bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] | None,
+    ) -> AxisAlignedBoundingBox:
+        """Return the per-row or default local bbox for *obj*, moved to *device*."""
+        if bboxes_per_row is not None and obj in bboxes_per_row:
+            return bboxes_per_row[obj].to(device)
+        return obj.get_bounding_box().to(device)
+
+    def _compute_total_loss(
+        self,
+        state: RelationSolverState,
+        debug: bool = False,
+        bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
+    ) -> torch.Tensor:
         """Compute total loss from all relations using registered strategies.
 
         Args:
             state: Current optimization state with object positions.
             debug: If True, print detailed loss breakdown.
+            bboxes_per_row: Optional per-row bboxes keyed by object. When
+                provided the bbox ``min_point`` / ``max_point`` have shape
+                ``(batch, 3)`` instead of ``(1, 3)``, enabling heterogeneous
+                object placement.
 
         Returns:
             Scalar loss tensor (mean over environments).
@@ -85,13 +106,14 @@ def _compute_total_loss(self, state: RelationSolverState, debug: bool = False) -
             for relation in obj.get_spatial_relations():
                 child_pos = state.get_position(obj)
                 strategy = self._get_strategy(relation)
+                child_bbox = self._get_bbox(obj, device, bboxes_per_row)
 
                 # Handle unary relations (no parent)
                 if isinstance(relation, UnaryRelation):
                     loss = strategy.compute_loss(
                         relation=relation,
                         child_pos=child_pos,
-                        child_bbox=obj.get_bounding_box().to(device),
+                        child_bbox=child_bbox,
                     )
                     if debug:
                         _print_unary_relation_debug(obj, relation, child_pos[0], loss.mean())
@@ -102,11 +124,12 @@ def _compute_total_loss(self, state: RelationSolverState, debug: bool = False) -
                         parent_world_bbox = parent.get_world_bounding_box().to(device)
                     else:
                         parent_pos = state.get_position(parent)
-                        parent_world_bbox = parent.get_bounding_box().to(device).translated(parent_pos)
+                        parent_bbox = self._get_bbox(parent, device, bboxes_per_row)
+                        parent_world_bbox = parent_bbox.translated(parent_pos)
                     loss = strategy.compute_loss(
                         relation=relation,
                         child_pos=child_pos,
-                        child_bbox=obj.get_bounding_box().to(device),
+                        child_bbox=child_bbox,
                         parent_world_bbox=parent_world_bbox,
                     )
                     if debug:
@@ -198,6 +221,7 @@ def solve(
         self,
         objects: list[ObjectBase],
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]],
+        bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> list[dict[ObjectBase, tuple[float, float, float]]]:
         """Solve for optimal positions of all objects.
 
@@ -206,6 +230,11 @@ def solve(
                 marked with IsAnchor() which serves as a fixed reference.
             initial_positions: List of dicts (one per env). Use a single-element list
                 for single-env placement.
+            bboxes_per_row: Optional per-row bounding boxes keyed by object.
+                When provided, each ``AxisAlignedBoundingBox`` has shape
+                ``(batch, 3)`` so different batch rows can use different
+                geometry (heterogeneous placement). If ``None``, every row
+                uses the object's default ``get_bounding_box()``.
 
         Returns:
             List of dicts (one per env) mapping objects to their solved (x, y, z) positions.
@@ -235,7 +264,7 @@ def solve(
         # Compute initial loss so _last_loss_per_env is always populated
         # (needed even when max_iters=0, e.g. tests that only check init positions).
         with torch.no_grad():
-            self._compute_total_loss(state)
+            self._compute_total_loss(state, bboxes_per_row=bboxes_per_row)
 
         # Optimization loop
         loss_history = []
@@ -248,7 +277,7 @@ def solve(
                 position_history.append(state.get_all_positions_snapshot())
 
             # Compute total loss
-            loss = self._compute_total_loss(state)
+            loss = self._compute_total_loss(state, bboxes_per_row=bboxes_per_row)
             loss_history.append(loss.item())
 
             # Backprop and update (only optimizable positions will update)
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
new file mode 100644
index 000000000..3659b30c8
--- /dev/null
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -0,0 +1,266 @@
+# 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
+
+"""Tests for heterogeneous object placement with per-env bounding boxes."""
+
+import torch
+
+from isaaclab_arena.assets.dummy_object import DummyObject
+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.relation_solver import RelationSolver
+from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
+from isaaclab_arena.relations.relation_solver_state import RelationSolverState
+from isaaclab_arena.relations.relations import IsAnchor, NoCollision, On
+from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+from isaaclab_arena.utils.pose import Pose
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+class HeterogeneousDummyObject(DummyObject):
+    """DummyObject that provides different bounding boxes per environment.
+
+    Used to exercise the heterogeneous placement path without requiring
+    RigidObjectSet's USD machinery.
+    """
+
+    def __init__(self, name: str, bboxes: list[AxisAlignedBoundingBox], **kwargs):
+        super().__init__(name=name, bounding_box=bboxes[0], **kwargs)
+        self._per_env_bboxes = bboxes
+        self.heterogeneous_bbox = True
+
+    def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
+        n_variants = len(self._per_env_bboxes)
+        indices = [i % n_variants for i in range(num_envs)]
+        min_pts = torch.stack([self._per_env_bboxes[idx].min_point[0] for idx in indices])
+        max_pts = torch.stack([self._per_env_bboxes[idx].max_point[0] for idx in indices])
+        return AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts)
+
+
+def _make_desk() -> DummyObject:
+    desk = DummyObject(
+        name="desk",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)),
+    )
+    desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)))
+    desk.add_relation(IsAnchor())
+    return desk
+
+
+# ---------------------------------------------------------------------------
+# ObjectBase.get_bounding_box_per_env
+# ---------------------------------------------------------------------------
+
+
+def test_dummy_object_bbox_per_env_expands_single():
+    """Default get_bounding_box_per_env should repeat the single bbox."""
+
+    obj = DummyObject(
+        name="box",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)),
+    )
+
+    per_env = obj.get_bounding_box_per_env(4)
+    assert per_env.min_point.shape == (4, 3)
+    assert per_env.max_point.shape == (4, 3)
+    assert torch.allclose(per_env.min_point[0], per_env.min_point[3])
+
+
+def test_heterogeneous_dummy_returns_different_bboxes():
+    """HeterogeneousDummyObject should cycle through its member bboxes."""
+
+    small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
+    large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3))
+    obj = HeterogeneousDummyObject(name="set", bboxes=[small, large])
+
+    per_env = obj.get_bounding_box_per_env(4)
+    assert per_env.max_point.shape == (4, 3)
+    # env 0 and 2 should use small; env 1 and 3 should use large
+    assert torch.allclose(per_env.max_point[0], torch.tensor([0.1, 0.1, 0.1]))
+    assert torch.allclose(per_env.max_point[1], torch.tensor([0.3, 0.3, 0.3]))
+
+
+# ---------------------------------------------------------------------------
+# Solver with per-row bboxes
+# ---------------------------------------------------------------------------
+
+
+def test_solver_accepts_per_row_bboxes():
+    """Solver should accept bboxes_per_row and produce valid results."""
+
+    desk = _make_desk()
+    box = DummyObject(
+        name="box",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)),
+    )
+    box.add_relation(On(desk, clearance_m=0.01))
+
+    objects = [desk, box]
+    batch_size = 4
+
+    initial_positions = [
+        {desk: (0.0, 0.0, 0.0), box: (0.5, 0.5, 0.11)} for _ in range(batch_size)
+    ]
+
+    # Create per-row bboxes with varying sizes across the batch.
+    min_pts = torch.zeros(batch_size, 3)
+    max_pts = torch.stack([torch.tensor([0.1 + 0.05 * i, 0.1 + 0.05 * i, 0.2]) for i in range(batch_size)])
+    per_row_bbox = AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts)
+
+    solver_params = RelationSolverParams(max_iters=100, convergence_threshold=1e-3, verbose=False)
+    solver = RelationSolver(params=solver_params)
+    result = solver.solve(objects, initial_positions, bboxes_per_row={box: per_row_bbox})
+
+    assert len(result) == batch_size
+    for pos_dict in result:
+        assert box in pos_dict
+        assert desk in pos_dict
+
+
+# ---------------------------------------------------------------------------
+# ObjectPlacer heterogeneous path
+# ---------------------------------------------------------------------------
+
+
+def test_placer_heterogeneous_produces_per_env_results():
+    """Placer should detect heterogeneous objects and solve per-env."""
+
+    desk = _make_desk()
+
+    small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
+    large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3))
+    hetero_box = HeterogeneousDummyObject(name="hetero_box", bboxes=[small, large])
+    hetero_box.add_relation(On(desk, clearance_m=0.01))
+
+    objects = [desk, hetero_box]
+    num_envs = 4
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    params = ObjectPlacerParams(
+        solver_params=solver_params,
+        apply_positions_to_objects=False,
+        placement_seed=None,
+    )
+
+    placer = ObjectPlacer(params=params)
+    result = placer.place(objects, num_envs=num_envs, result_per_env=True)
+
+    assert isinstance(result, MultiEnvPlacementResult)
+    assert len(result.results) == num_envs
+    for r in result.results:
+        assert hetero_box in r.positions
+
+
+def test_placer_heterogeneous_z_height_matches_variant():
+    """Objects should be placed at z-height matching their env's variant bbox."""
+
+    desk = _make_desk()
+
+    # "tall" variant: height 0.4 → bottom at z ≈ 0.11 (desk top 0.1 + clearance 0.01)
+    tall = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.4))
+    # "short" variant: height 0.1 → bottom at z ≈ 0.11 (same clearance)
+    short = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.1))
+    hetero = HeterogeneousDummyObject(name="hetero", bboxes=[tall, short])
+    hetero.add_relation(On(desk, clearance_m=0.01))
+
+    objects = [desk, hetero]
+    num_envs = 2
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    params = ObjectPlacerParams(
+        solver_params=solver_params,
+        apply_positions_to_objects=False,
+        placement_seed=42,
+    )
+
+    placer = ObjectPlacer(params=params)
+    result = placer.place(objects, num_envs=num_envs, result_per_env=True)
+
+    assert isinstance(result, MultiEnvPlacementResult)
+    # Both envs should have solved z near the desk top + clearance (0.11).
+    # The On loss targets: z = parent_top + clearance - child_min_z = 0.1 + 0.01 - 0.0 = 0.11
+    for env_idx, r in enumerate(result.results):
+        z = r.positions[hetero][2]
+        assert abs(z - 0.11) < 0.05, f"Env {env_idx}: z={z:.4f}, expected ~0.11"
+
+
+def test_mixed_heterogeneous_and_homogeneous_placement():
+    """Mixed scene: heterogeneous A (RigidObjectSet-like) + homogeneous X (plain Object).
+
+    Both sit On(desk) with NoCollision between them. The solver must produce
+    valid, non-overlapping placements in every env even though A has different
+    bboxes per env while X is identical everywhere.
+    """
+
+    desk = _make_desk()
+
+    # A: heterogeneous — small variant in even envs, large in odd envs.
+    small_a = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
+    large_a = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.25, 0.25, 0.25))
+    obj_a = HeterogeneousDummyObject(name="A", bboxes=[small_a, large_a])
+    obj_a.add_relation(On(desk, clearance_m=0.01))
+
+    # X: homogeneous — same bbox in all envs.
+    obj_x = DummyObject(
+        name="X",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.15)),
+    )
+    obj_x.add_relation(On(desk, clearance_m=0.01))
+
+    # NoCollision between A and X (both directions).
+    obj_a.add_relation(NoCollision(obj_x))
+    obj_x.add_relation(NoCollision(obj_a))
+
+    objects = [desk, obj_a, obj_x]
+    num_envs = 4
+
+    solver_params = RelationSolverParams(max_iters=300, convergence_threshold=1e-3, verbose=False)
+    params = ObjectPlacerParams(
+        solver_params=solver_params,
+        apply_positions_to_objects=False,
+        placement_seed=42,
+    )
+
+    placer = ObjectPlacer(params=params)
+    result = placer.place(objects, num_envs=num_envs, result_per_env=True)
+
+    assert isinstance(result, MultiEnvPlacementResult)
+    assert len(result.results) == num_envs
+
+    for env_idx, r in enumerate(result.results):
+        assert obj_a in r.positions and obj_x in r.positions
+        # Verify z-height is near desk top + clearance for both objects.
+        for obj in (obj_a, obj_x):
+            z = r.positions[obj][2]
+            assert abs(z - 0.11) < 0.05, f"Env {env_idx}, {obj.name}: z={z:.4f}, expected ~0.11"
+
+
+def test_homogeneous_path_unchanged():
+    """When no heterogeneous objects exist, the homogeneous path is used."""
+
+    desk = _make_desk()
+    box = DummyObject(
+        name="box",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)),
+    )
+    box.add_relation(On(desk, clearance_m=0.01))
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    params = ObjectPlacerParams(
+        solver_params=solver_params,
+        apply_positions_to_objects=False,
+        placement_seed=None,
+    )
+
+    placer = ObjectPlacer(params=params)
+    result = placer.place([desk, box], num_envs=2, result_per_env=True)
+
+    assert isinstance(result, MultiEnvPlacementResult)
+    assert len(result.results) == 2

From f6d5594c585fff17810ac513cd3aaa835f80253c Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 15 Apr 2026 14:06:57 -0700
Subject: [PATCH 02/46] rebase

---
 isaaclab_arena/assets/object_set.py           | 11 ++++-
 isaaclab_arena/relations/object_placer.py     | 47 +++++++++++++++----
 isaaclab_arena/relations/relation_solver.py   | 16 +++++--
 .../tests/test_heterogeneous_placement.py     | 22 +++++----
 4 files changed, 72 insertions(+), 24 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 6df45acf6..19c94070f 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -93,7 +93,7 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """
         return max(self.objects, key=lambda obj: obj.get_bounding_box().size[0, 2].item()).get_bounding_box()
 
-    def get_variant_indices(self, num_envs: int) -> list[int]:
+    def get_variant_indices(self, num_envs: int, seed: int | None = None) -> list[int]:
         """Return which member object index is assigned to each environment.
 
         When ``random_choice`` is False the mapping is round-robin
@@ -103,6 +103,9 @@ def get_variant_indices(self, num_envs: int) -> list[int]:
 
         Args:
             num_envs: Number of environments.
+            seed: Optional RNG seed for reproducible variant assignment
+                when ``random_choice`` is True. If None, uses the global
+                torch RNG.
 
         Returns:
             List of length ``num_envs`` with indices into ``self.objects``.
@@ -112,7 +115,11 @@ def get_variant_indices(self, num_envs: int) -> list[int]:
             return [i % n for i in range(num_envs)]
 
         if not hasattr(self, "_cached_variant_indices") or len(self._cached_variant_indices) != num_envs:
-            self._cached_variant_indices = [int(i) for i in torch.randint(n, (num_envs,)).tolist()]
+            generator = None
+            if seed is not None:
+                generator = torch.Generator()
+                generator.manual_seed(seed)
+            self._cached_variant_indices = [int(i) for i in torch.randint(n, (num_envs,), generator=generator).tolist()]
         return self._cached_variant_indices
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index 7eaa7d82f..464c2a007 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -253,11 +253,19 @@ def _place_heterogeneous(
         results: list[PlacementResult] = []
         for env_idx in range(num_envs):
             start = env_idx * max_attempts
+            # Slice single-env bboxes for validation of this env's candidates.
+            env_bbox_overrides = {
+                obj: AxisAlignedBoundingBox(
+                    min_point=env_bboxes[obj].min_point[env_idx : env_idx + 1],
+                    max_point=env_bboxes[obj].max_point[env_idx : env_idx + 1],
+                )
+                for obj in objects
+            }
             env_candidates = [
                 PlacementCandidate(
                     all_losses[start + j],
                     all_positions[start + j],
-                    self._validate_placement(all_positions[start + j]),
+                    self._validate_placement(all_positions[start + j], bbox_overrides=env_bbox_overrides),
                 )
                 for j in range(max_attempts)
             ]
@@ -415,12 +423,18 @@ def _sample_axis_position(
     def _validate_on_relations(
         self,
         positions: dict[ObjectBase, tuple[float, float, float]],
+        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> bool:
         """Validate each On relation; logic matches OnLossStrategy (relation_loss_strategies.py).
 
         1. X: child's footprint entirely within parent's X extent.
         2. Y: child's footprint entirely within parent's Y extent.
         3. Z: child_bottom in (parent_top, parent_top+clearance_m], within on_relation_z_tolerance_m.
+
+        Args:
+            positions: Solved positions for each object.
+            bbox_overrides: Optional per-object bbox overrides (single-env, shape ``(1, 3)``).
+                Used by heterogeneous placement to supply the correct variant bbox.
         """
         for obj in positions:
             for rel in obj.get_relations():
@@ -429,8 +443,12 @@ def _validate_on_relations(
                 parent = rel.parent
                 if parent not in positions:
                     continue
-                child_world = obj.get_bounding_box().translated(positions[obj])
-                parent_world = parent.get_bounding_box().translated(positions[parent])
+                child_bbox = bbox_overrides[obj] if bbox_overrides and obj in bbox_overrides else obj.get_bounding_box()
+                parent_bbox = (
+                    bbox_overrides[parent] if bbox_overrides and parent in bbox_overrides else parent.get_bounding_box()
+                )
+                child_world = child_bbox.translated(positions[obj])
+                parent_world = parent_bbox.translated(positions[parent])
                 # 1 & 2: Same as OnLossStrategy X/Y band (child's footprint within parent).
                 if (
                     child_world.min_point[0, 0] < parent_world.min_point[0, 0]
@@ -442,8 +460,8 @@ def _validate_on_relations(
                         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[0, 2].item()
-                child_local_bottom_z: float = obj.get_bounding_box().min_point[0, 2].item()
+                parent_local_top_z: float = parent_bbox.max_point[0, 2].item()
+                child_local_bottom_z: float = child_bbox.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]
@@ -457,6 +475,7 @@ def _validate_on_relations(
     def _validate_no_overlap(
         self,
         positions: dict[ObjectBase, tuple[float, float, float]],
+        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> bool:
         """Validate that no two objects overlap in 3D (axis-aligned bbox with margin).
 
@@ -464,6 +483,11 @@ def _validate_no_overlap(
         The margin is derived from the solver's clearance_m parameter (with a
         small float tolerance subtracted to avoid rejecting solutions that are
         within solver residual).
+
+        Args:
+            positions: Solved positions for each object.
+            bbox_overrides: Optional per-object bbox overrides (single-env, shape ``(1, 3)``).
+                Used by heterogeneous placement to supply the correct variant bbox.
         """
         # Build set of On-related pairs to skip (child, parent) and (parent, child).
         on_pairs: set[tuple] = set()
@@ -490,8 +514,10 @@ def _validate_no_overlap(
                 if (id(a), id(b)) in on_pairs:
                     continue
 
-                a_world = a.get_bounding_box().translated(positions[a])
-                b_world = b.get_bounding_box().translated(positions[b])
+                a_bbox = bbox_overrides[a] if bbox_overrides and a in bbox_overrides else a.get_bounding_box()
+                b_bbox = bbox_overrides[b] if bbox_overrides and b in bbox_overrides else b.get_bounding_box()
+                a_world = a_bbox.translated(positions[a])
+                b_world = b_bbox.translated(positions[b])
 
                 if a_world.overlaps(b_world, margin=margin).item():
                     if self.params.verbose:
@@ -502,16 +528,21 @@ def _validate_no_overlap(
     def _validate_placement(
         self,
         positions: dict[ObjectBase, tuple[float, float, float]],
+        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> bool:
         """Validate that no two objects overlap in 3D and On relations are satisfied.
 
         Args:
             positions: Dictionary mapping objects to their solved (x, y, z) positions.
+            bbox_overrides: Optional per-object bbox overrides (single-env, shape ``(1, 3)``).
+                Used by heterogeneous placement to supply the correct variant bbox.
 
         Returns:
             True if no overlaps exist and On relations hold, False otherwise.
         """
-        return self._validate_no_overlap(positions) and self._validate_on_relations(positions)
+        return self._validate_no_overlap(positions, bbox_overrides) and self._validate_on_relations(
+            positions, bbox_overrides
+        )
 
     def _apply_positions(
         self,
diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index ebc042b14..844c83a4f 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -15,7 +15,7 @@
 )
 from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
 from isaaclab_arena.relations.relation_solver_state import RelationSolverState
-from isaaclab_arena.relations.relations import AtPosition, Relation, RelationBase, UnaryRelation
+from isaaclab_arena.relations.relations import Relation, RelationBase, UnaryRelation
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
@@ -141,12 +141,17 @@ def _compute_total_loss(
                 total_loss = total_loss + loss
 
         # Add built-in no-overlap loss between all object pairs
-        total_loss = total_loss + self._compute_no_overlap_loss(state, debug)
+        total_loss = total_loss + self._compute_no_overlap_loss(state, debug, bboxes_per_row=bboxes_per_row)
 
         self._last_loss_per_env = total_loss.detach().clone()
         return total_loss.mean()
 
-    def _compute_no_overlap_loss(self, state: RelationSolverState, debug: bool = False) -> torch.Tensor:
+    def _compute_no_overlap_loss(
+        self,
+        state: RelationSolverState,
+        debug: bool = False,
+        bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
+    ) -> torch.Tensor:
         """Compute pairwise no-overlap loss for all non-anchor objects against all other objects.
 
         Each unique pair is evaluated twice (once per direction):
@@ -157,6 +162,7 @@ def _compute_no_overlap_loss(self, state: RelationSolverState, debug: bool = Fal
         Args:
             state: Current optimization state with object positions.
             debug: If True, print detailed loss breakdown.
+            bboxes_per_row: Optional per-row bboxes for heterogeneous placement.
 
         Returns:
             Per-environment loss tensor of shape (batch_size,).
@@ -169,7 +175,7 @@ def _compute_no_overlap_loss(self, state: RelationSolverState, debug: bool = Fal
 
         for i, child in enumerate(non_anchor_objects):
             child_pos = state.get_position(child)
-            child_bbox = child.get_bounding_box().to(device)
+            child_bbox = self._get_bbox(child, device, bboxes_per_row)
 
             # Against all anchors
             for anchor in anchor_objects:
@@ -188,7 +194,7 @@ def _compute_no_overlap_loss(self, state: RelationSolverState, debug: bool = Fal
             for j in range(i + 1, len(non_anchor_objects)):
                 other = non_anchor_objects[j]
                 other_pos = state.get_position(other)
-                other_bbox = other.get_bounding_box().to(device)
+                other_bbox = self._get_bbox(other, device, bboxes_per_row)
 
                 # Forward: gradient flows to child (object i)
                 other_world_bbox = other_bbox.translated(other_pos.detach())
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 3659b30c8..20335180b 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -13,12 +13,10 @@
 from isaaclab_arena.relations.placement_result import MultiEnvPlacementResult
 from isaaclab_arena.relations.relation_solver import RelationSolver
 from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
-from isaaclab_arena.relations.relation_solver_state import RelationSolverState
-from isaaclab_arena.relations.relations import IsAnchor, NoCollision, On
+from isaaclab_arena.relations.relations import IsAnchor, On
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 from isaaclab_arena.utils.pose import Pose
 
-
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -105,9 +103,7 @@ def test_solver_accepts_per_row_bboxes():
     objects = [desk, box]
     batch_size = 4
 
-    initial_positions = [
-        {desk: (0.0, 0.0, 0.0), box: (0.5, 0.5, 0.11)} for _ in range(batch_size)
-    ]
+    initial_positions = [{desk: (0.0, 0.0, 0.0), box: (0.5, 0.5, 0.11)} for _ in range(batch_size)]
 
     # Create per-row bboxes with varying sizes across the batch.
     min_pts = torch.zeros(batch_size, 3)
@@ -214,9 +210,7 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
     )
     obj_x.add_relation(On(desk, clearance_m=0.01))
 
-    # NoCollision between A and X (both directions).
-    obj_a.add_relation(NoCollision(obj_x))
-    obj_x.add_relation(NoCollision(obj_a))
+    # No-overlap is handled automatically by the solver's built-in clearance.
 
     objects = [desk, obj_a, obj_x]
     num_envs = 4
@@ -241,6 +235,16 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
             z = r.positions[obj][2]
             assert abs(z - 0.11) < 0.05, f"Env {env_idx}, {obj.name}: z={z:.4f}, expected ~0.11"
 
+        # Verify XY non-overlap using each env's actual variant bbox.
+        variant_idx = env_idx % len([small_a, large_a])
+        a_bbox = [small_a, large_a][variant_idx]
+        x_bbox = obj_x.get_bounding_box()
+        a_world = a_bbox.translated(r.positions[obj_a])
+        x_world = x_bbox.translated(r.positions[obj_x])
+        assert not a_world.overlaps(
+            x_world
+        ).item(), f"Env {env_idx}: A and X bboxes overlap at positions A={r.positions[obj_a]}, X={r.positions[obj_x]}"
+
 
 def test_homogeneous_path_unchanged():
     """When no heterogeneous objects exist, the homogeneous path is used."""

From 6d25a3afaa845d5709a5bf41eeeee35cdfb651a0 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 21 Apr 2026 09:38:19 -0700
Subject: [PATCH 03/46] Add heterogeneous placement support to
 PooledObjectPlacer

Extends the existing PooledObjectPlacer with per-variant sub-pools
for heterogeneous objects (e.g. RigidObjectSet with varying bboxes).
Keeps the peer-reviewed API from placement-on-reset unchanged:
sample_without_replacement() gains an optional env_ids parameter,
sample_with_replacement() routes to variant-aware sampling internally.

Made-with: Cursor
---
 .../environments/relation_solver_interface.py |   1 +
 isaaclab_arena/relations/placement_events.py  |   2 +-
 .../relations/pooled_object_placer.py         | 207 ++++++++++++++++--
 3 files changed, 193 insertions(+), 17 deletions(-)

diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py
index 6f304a52f..8f6bf6bf5 100644
--- a/isaaclab_arena/environments/relation_solver_interface.py
+++ b/isaaclab_arena/environments/relation_solver_interface.py
@@ -59,6 +59,7 @@ def solve_and_apply_relation_placement(
         objects=objects,
         placer_params=placer_params,
         pool_size=num_envs * placer_params.min_unique_layouts_per_env,
+        num_envs=num_envs,
     )
     return _apply_relation_placement_result(
         objects=objects,
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index eac560c66..ff0545b41 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -47,7 +47,7 @@ def solve_and_place_objects(
         return
 
     num_reset_envs = len(env_ids)
-    layouts_per_env = placement_pool.sample_without_replacement(num_reset_envs)
+    layouts_per_env = placement_pool.sample_without_replacement(num_reset_envs, env_ids=env_ids)
 
     anchor_objects_set = set(get_anchor_objects(objects))
     rotations = {obj: get_rotation_xyzw(obj) for obj in objects if obj not in anchor_objects_set}
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 5751c7a5d..befd68ce3 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -24,8 +24,22 @@ class PooledObjectPlacer:
     ``pool_size``, keeping only those that pass validation. The pool is refilled
     automatically when consumed layouts run out.
 
+    **Homogeneous mode** (default): all objects have the same geometry in
+    every environment.  Layouts are stored in a single flat list and any
+    layout can serve any environment.
+
+    **Heterogeneous mode** (activated when any object has
+    ``heterogeneous_bbox = True``, e.g. ``RigidObjectSet``): each layout
+    is tied to a specific *variant index* (``env_idx % num_variants``).
+    Layouts are bucketed into per-variant sub-pools so that
+    :meth:`sample_without_replacement` and :meth:`sample_with_replacement`
+    always return a layout that matches the requesting environment's
+    variant geometry.
+
     * :meth:`sample_without_replacement` — returns the next *count* layouts
-      sequentially.  Auto-refills when exhausted.
+      sequentially.  Auto-refills when exhausted.  In heterogeneous mode,
+      pass ``env_ids`` so each environment receives a layout matching its
+      variant.
     * :meth:`sample_with_replacement` — picks *count* layouts at random
       (non-consuming).  Used for static initial positions.
 
@@ -33,6 +47,8 @@ class PooledObjectPlacer:
         objects: All objects (including anchors) participating in relation solving.
         placer_params: Parameters forwarded to ``ObjectPlacer`` for the batched solve.
         pool_size: Number of layouts to solve per batch.
+        num_envs: Total number of simulation environments.  Required for
+            heterogeneous placement so variant indices can be resolved.
     """
 
     def __init__(
@@ -40,23 +56,69 @@ def __init__(
         objects: list[ObjectBase],
         placer_params: ObjectPlacerParams,
         pool_size: int = 100,
+        num_envs: int | None = None,
     ) -> None:
         if pool_size < 1:
             raise ValueError(f"pool_size must be >= 1, got {pool_size}")
 
-        self._objects = objects
+        self._objects = list(objects)
         self._placer = ObjectPlacer(params=placer_params)
         self._pool_size = pool_size
-        self._layouts: list[PlacementResult] = []
-        self._next_idx: int = 0
+        self._heterogeneous = any(getattr(obj, "heterogeneous_bbox", False) for obj in objects)
 
-        # Pre-solve the initial batch (runs the gradient solver, no simulation is needed).
-        self._solve_and_store(pool_size)
-        if not self._layouts:
-            raise RuntimeError(
-                f"Pooled object placer failed to produce any valid layouts from {pool_size} attempts. "
-                "Check object relations and constraints."
-            )
+        if self._heterogeneous:
+            assert (
+                num_envs is not None
+            ), "num_envs is required for heterogeneous placement pools so variant indices can be resolved."
+            self._num_envs = num_envs
+            self._num_variants = self._detect_num_variants(objects)
+
+            self._variant_layouts: dict[int, list[PlacementResult]] = {v: [] for v in range(self._num_variants)}
+            self._variant_next_idx: dict[int, int] = {v: 0 for v in range(self._num_variants)}
+
+            self._solve_and_store_heterogeneous(pool_size)
+            for v in range(self._num_variants):
+                if not self._variant_layouts[v]:
+                    raise RuntimeError(
+                        f"Placement pool failed to produce any valid layouts for variant {v} "
+                        f"from {pool_size} attempts. Check object relations and constraints."
+                    )
+        else:
+            self._layouts: list[PlacementResult] = []
+            self._next_idx: int = 0
+
+            # Pre-solve the initial batch (runs the gradient solver, no simulation is needed).
+            self._solve_and_store(pool_size)
+            if not self._layouts:
+                raise RuntimeError(
+                    f"Pooled object placer failed to produce any valid layouts from {pool_size} attempts. "
+                    "Check object relations and constraints."
+                )
+
+    @property
+    def is_heterogeneous(self) -> bool:
+        """Whether this pool operates in heterogeneous (per-variant) mode."""
+        return self._heterogeneous
+
+    # ------------------------------------------------------------------
+    # Variant helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _detect_num_variants(objects: list[ObjectBase]) -> int:
+        """Return the number of unique object variants across heterogeneous objects."""
+        for obj in objects:
+            if getattr(obj, "heterogeneous_bbox", False):
+                return len(obj.objects)  # type: ignore[attr-defined]
+        return 1
+
+    def _variant_for_env(self, env_id: int) -> int:
+        """Map an environment index to its variant index."""
+        return env_id % self._num_variants
+
+    # ------------------------------------------------------------------
+    # Homogeneous (flat pool) internals
+    # ------------------------------------------------------------------
 
     def _compact(self) -> None:
         """Drop consumed layouts and reset the read index to free memory."""
@@ -72,8 +134,6 @@ def _solve_and_store(self, num_layouts: int) -> None:
         """
         self._compact()
 
-        # place() runs: random init → gradient solve → validate → rank.
-        # It returns up to num_layouts results; some may fail validation.
         with torch.inference_mode(False):
             result = self._placer.place(self._objects, num_envs=num_layouts, result_per_env=True)
 
@@ -104,14 +164,73 @@ def _refill_pool_via_solve_if_required(self, count: int) -> None:
                 "The solver is not producing enough valid placements."
             )
 
-    def sample_without_replacement(self, count: int) -> list[PlacementResult]:
+    # ------------------------------------------------------------------
+    # Heterogeneous (per-variant sub-pool) internals
+    # ------------------------------------------------------------------
+
+    def _compact_variant(self, variant: int) -> None:
+        """Drop consumed layouts for a single variant and reset its read index."""
+        idx = self._variant_next_idx[variant]
+        self._variant_layouts[variant] = self._variant_layouts[variant][idx:]
+        self._variant_next_idx[variant] = 0
+
+    def _solve_and_store_heterogeneous(self, num_layouts: int) -> None:
+        """Solve layouts and bucket valid results by variant index.
+
+        Each result ``i`` from the solver corresponds to variant
+        ``i % num_variants`` because ``get_bounding_box_per_env`` assigns
+        variants in round-robin order.
+        """
+        for v in range(self._num_variants):
+            self._compact_variant(v)
+
+        with torch.inference_mode(False):
+            result = self._placer.place(self._objects, num_envs=num_layouts, result_per_env=True)
+
+        all_results = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
+
+        total_valid = 0
+        for i, r in enumerate(all_results):
+            if r.success:
+                variant = i % self._num_variants
+                self._variant_layouts[variant].append(r)
+                total_valid += 1
+
+        if total_valid < num_layouts:
+            print(
+                f"Placement pool (heterogeneous): solved {num_layouts} candidates,"
+                f" {total_valid} valid, {num_layouts - total_valid} failed validation"
+            )
+
+        for v in range(self._num_variants):
+            random.shuffle(self._variant_layouts[v])
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def sample_without_replacement(
+        self, count: int, env_ids: list[int] | torch.Tensor | None = None
+    ) -> list[PlacementResult]:
         """Return the next *count* layouts sequentially (without replacement).
 
-        Auto-refills the pool when there are not enough layouts ahead of the read index.
+        Auto-refills the pool when there are not enough layouts ahead of the
+        read index.
+
+        In **heterogeneous mode** ``env_ids`` must be provided so each
+        environment receives a layout matching its variant geometry.
+
+        Args:
+            count: Number of layouts to return.
+            env_ids: Environment indices being reset.  Required when the
+                pool is heterogeneous; ignored otherwise.
 
         Raises:
             RuntimeError: If the pool cannot provide *count* layouts after refilling.
         """
+        if self._heterogeneous:
+            return self._sample_without_replacement_heterogeneous(count, env_ids)
+
         self._refill_pool_via_solve_if_required(count)
 
         start = self._next_idx
@@ -119,17 +238,73 @@ def sample_without_replacement(self, count: int) -> list[PlacementResult]:
         layouts = self._layouts[start : self._next_idx]
         return layouts
 
+    def _sample_without_replacement_heterogeneous(
+        self, count: int, env_ids: list[int] | torch.Tensor | None
+    ) -> list[PlacementResult]:
+        assert env_ids is not None, "env_ids must be provided for heterogeneous placement pools."
+
+        if isinstance(env_ids, torch.Tensor):
+            ids: list[int] = [int(x) for x in env_ids]
+        else:
+            ids = list(env_ids)
+        assert len(ids) == count
+
+        variant_demand: dict[int, int] = {}
+        for eid in ids:
+            v = self._variant_for_env(eid)
+            variant_demand[v] = variant_demand.get(v, 0) + 1
+
+        for v, demand in variant_demand.items():
+            avail = len(self._variant_layouts[v]) - self._variant_next_idx[v]
+            if avail < demand:
+                refill = max(self._pool_size, demand * self._num_variants)
+                self._solve_and_store_heterogeneous(refill)
+
+        results: list[PlacementResult] = []
+        for eid in ids:
+            v = self._variant_for_env(eid)
+            idx = self._variant_next_idx[v]
+            if idx >= len(self._variant_layouts[v]):
+                raise RuntimeError(
+                    f"Placement pool: variant {v} has no more valid layouts "
+                    f"(needed for env {eid}). The solver is not producing enough valid placements."
+                )
+            results.append(self._variant_layouts[v][idx])
+            self._variant_next_idx[v] = idx + 1
+
+        return results
+
     def sample_with_replacement(self, count: int) -> list[PlacementResult]:
         """Pick *count* layouts at random with replacement (non-consuming).
 
+        In **heterogeneous mode**, each position ``i`` in the returned
+        list corresponds to env ``i`` and is drawn from the sub-pool
+        matching that env's variant (``i % num_variants``).
+
         Used by ``resolve_on_reset=False`` to assign initial positions
         that persist across resets.
         """
+        if self._heterogeneous:
+            return self._sample_with_replacement_heterogeneous(count)
         return random.choices(self._layouts, k=count)
 
+    def _sample_with_replacement_heterogeneous(self, count: int) -> list[PlacementResult]:
+        results: list[PlacementResult] = []
+        for env_idx in range(count):
+            v = self._variant_for_env(env_idx)
+            pool = self._variant_layouts[v]
+            assert pool, f"Variant {v} has no valid layouts to sample from."
+            results.append(random.choice(pool))
+        return results
+
     @property
     def remaining(self) -> int:
-        """Number of layouts not yet consumed by :meth:`sample_without_replacement`."""
+        """Number of layouts not yet consumed by :meth:`sample_without_replacement`.
+
+        For heterogeneous pools, returns the minimum across all variants.
+        """
+        if self._heterogeneous:
+            return min(len(self._variant_layouts[v]) - self._variant_next_idx[v] for v in range(self._num_variants))
         return len(self._layouts) - self._next_idx
 
     @property

From 4a7c1ea3a8e00af4a4dbae17f6d32daff69fdb6f Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Fri, 24 Apr 2026 10:41:44 -0700
Subject: [PATCH 04/46] Add assert, fallback, and pool tests from vic change

- Assert random_choice=False when heterogeneous_bbox=True in RigidObjectSet
- Add fallback in heterogeneous pool when no candidates pass validation
- Add PooledObjectPlacer heterogeneous mode tests
- Remove unnecessary list() defensive copy

Made-with: Cursor
---
 isaaclab_arena/assets/object_set.py           |  6 ++
 .../relations/pooled_object_placer.py         |  9 +-
 .../tests/test_heterogeneous_placement.py     | 99 +++++++++++++++++++
 3 files changed, 111 insertions(+), 3 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 19c94070f..5cc1e1643 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -71,6 +71,12 @@ def __init__(
         self.random_choice = random_choice
         self.heterogeneous_bbox: bool = len(objects) > 1
 
+        assert not (self.heterogeneous_bbox and self.random_choice), (
+            f"RigidObjectSet '{name}': random_choice=True is not supported with heterogeneous "
+            "placement (len(objects) > 1). The placement pool assumes round-robin variant "
+            "assignment (env_idx % num_variants) which conflicts with random spawning order."
+        )
+
         # Set default prim_path if not provided
         if prim_path is None:
             prim_path = f"{{ENV_REGEX_NS}}/{name}"
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index befd68ce3..4a8f5e22b 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -61,7 +61,7 @@ def __init__(
         if pool_size < 1:
             raise ValueError(f"pool_size must be >= 1, got {pool_size}")
 
-        self._objects = list(objects)
+        self._objects = objects
         self._placer = ObjectPlacer(params=placer_params)
         self._pool_size = pool_size
         self._heterogeneous = any(getattr(obj, "heterogeneous_bbox", False) for obj in objects)
@@ -202,8 +202,11 @@ def _solve_and_store_heterogeneous(self, num_layouts: int) -> None:
                 f" {total_valid} valid, {num_layouts - total_valid} failed validation"
             )
 
-        for v in range(self._num_variants):
-            random.shuffle(self._variant_layouts[v])
+        if total_valid == 0:
+            print("Warning: No candidates passed strict validation. Accepting best-loss layouts as fallback.")
+            for i, r in enumerate(all_results):
+                variant = i % self._num_variants
+                self._variant_layouts[variant].append(r)
 
     # ------------------------------------------------------------------
     # Public API
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 20335180b..ab49efc99 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -11,6 +11,7 @@
 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.pooled_object_placer import PooledObjectPlacer
 from isaaclab_arena.relations.relation_solver import RelationSolver
 from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
 from isaaclab_arena.relations.relations import IsAnchor, On
@@ -33,6 +34,7 @@ def __init__(self, name: str, bboxes: list[AxisAlignedBoundingBox], **kwargs):
         super().__init__(name=name, bounding_box=bboxes[0], **kwargs)
         self._per_env_bboxes = bboxes
         self.heterogeneous_bbox = True
+        self.objects = bboxes
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         n_variants = len(self._per_env_bboxes)
@@ -268,3 +270,100 @@ def test_homogeneous_path_unchanged():
 
     assert isinstance(result, MultiEnvPlacementResult)
     assert len(result.results) == 2
+
+
+# ---------------------------------------------------------------------------
+# PooledObjectPlacer heterogeneous mode
+# ---------------------------------------------------------------------------
+
+
+def _make_hetero_pool_objects():
+    """Create desk + heterogeneous box for pool tests."""
+    desk = _make_desk()
+    small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
+    large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3))
+    hetero = HeterogeneousDummyObject(name="hetero", bboxes=[small, large])
+    hetero.add_relation(On(desk, clearance_m=0.01))
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
+    return desk, hetero, placer_params
+
+
+def test_pooled_placer_heterogeneous_is_detected():
+    """PooledObjectPlacer should detect heterogeneous objects and create variant sub-pools."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
+
+    assert pool.is_heterogeneous
+    assert pool.remaining > 0
+
+
+def test_pooled_placer_heterogeneous_sample_without_replacement():
+    """sample_without_replacement with env_ids should return one layout per env from correct variant."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
+
+    env_ids = torch.tensor([0, 1, 2, 3])
+    draws = pool.sample_without_replacement(4, env_ids=env_ids)
+    assert len(draws) == 4
+    for d in draws:
+        assert hetero in d.positions
+
+
+def test_pooled_placer_heterogeneous_sample_without_replacement_requires_env_ids():
+    """Heterogeneous pool should assert when env_ids is not provided."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
+
+    try:
+        pool.sample_without_replacement(2, env_ids=None)
+        assert False, "Should have raised AssertionError"
+    except AssertionError:
+        pass
+
+
+def test_pooled_placer_heterogeneous_sample_with_replacement():
+    """sample_with_replacement should return per-variant layouts without consuming."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
+
+    initial_remaining = pool.remaining
+    samples = pool.sample_with_replacement(4)
+    assert len(samples) == 4
+    assert pool.remaining == initial_remaining, "sample_with_replacement should not consume layouts"
+
+
+def test_pooled_placer_heterogeneous_refill():
+    """Exhausting a variant sub-pool should trigger a refill."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=4, num_envs=2)
+
+    initial_remaining = pool.remaining
+
+    # Draw all layouts for env 0 (variant 0) and env 1 (variant 1)
+    env_ids = torch.tensor([0, 1] * initial_remaining)
+    pool.sample_without_replacement(len(env_ids), env_ids=env_ids)
+
+    # Pool should be exhausted now; request more to trigger refill
+    env_ids_more = torch.tensor([0, 1])
+    draws = pool.sample_without_replacement(2, env_ids=env_ids_more)
+    assert len(draws) == 2, "Pool should refill and return requested layouts"
+
+
+def test_pooled_placer_homogeneous_unaffected_by_num_envs():
+    """Homogeneous pool should work the same whether num_envs is passed or not."""
+    desk = _make_desk()
+    box = DummyObject(
+        name="box",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)),
+    )
+    box.add_relation(On(desk, clearance_m=0.01))
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
+
+    pool = PooledObjectPlacer(objects=[desk, box], placer_params=placer_params, pool_size=10, num_envs=4)
+    assert not pool.is_heterogeneous
+    draws = pool.sample_without_replacement(3)
+    assert len(draws) == 3

From b7f4a46324033025fab5fab34e305a5557a50ca0 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Fri, 24 Apr 2026 11:10:10 -0700
Subject: [PATCH 05/46] fix heter fallback

---
 .../relations/pooled_object_placer.py         | 193 +++++++++---------
 .../tests/test_heterogeneous_placement.py     | 119 +++++++++++
 2 files changed, 220 insertions(+), 92 deletions(-)

diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 4a8f5e22b..00507c03c 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -29,17 +29,14 @@ class PooledObjectPlacer:
     layout can serve any environment.
 
     **Heterogeneous mode** (activated when any object has
-    ``heterogeneous_bbox = True``, e.g. ``RigidObjectSet``): each layout
-    is tied to a specific *variant index* (``env_idx % num_variants``).
-    Layouts are bucketed into per-variant sub-pools so that
-    :meth:`sample_without_replacement` and :meth:`sample_with_replacement`
-    always return a layout that matches the requesting environment's
-    variant geometry.
+    ``heterogeneous_bbox = True``, e.g. ``RigidObjectSet``): each
+    environment has its own fixed set of object variants, assigned at
+    build time.  Layouts are stored per ``env_id`` so that resets always
+    return a layout solved for that environment's actual object geometry.
 
     * :meth:`sample_without_replacement` — returns the next *count* layouts
       sequentially.  Auto-refills when exhausted.  In heterogeneous mode,
-      pass ``env_ids`` so each environment receives a layout matching its
-      variant.
+      pass ``env_ids`` so each environment receives a matching layout.
     * :meth:`sample_with_replacement` — picks *count* layouts at random
       (non-consuming).  Used for static initial positions.
 
@@ -48,7 +45,7 @@ class PooledObjectPlacer:
         placer_params: Parameters forwarded to ``ObjectPlacer`` for the batched solve.
         pool_size: Number of layouts to solve per batch.
         num_envs: Total number of simulation environments.  Required for
-            heterogeneous placement so variant indices can be resolved.
+            heterogeneous placement so per-env pools can be created.
     """
 
     def __init__(
@@ -69,18 +66,17 @@ def __init__(
         if self._heterogeneous:
             assert (
                 num_envs is not None
-            ), "num_envs is required for heterogeneous placement pools so variant indices can be resolved."
+            ), "num_envs is required for heterogeneous placement so per-env pools can be created."
             self._num_envs = num_envs
-            self._num_variants = self._detect_num_variants(objects)
 
-            self._variant_layouts: dict[int, list[PlacementResult]] = {v: [] for v in range(self._num_variants)}
-            self._variant_next_idx: dict[int, int] = {v: 0 for v in range(self._num_variants)}
+            self._layout_pools: dict[int, list[PlacementResult]] = {env_id: [] for env_id in range(num_envs)}
+            self._layout_cursors: dict[int, int] = {env_id: 0 for env_id in range(num_envs)}
 
             self._solve_and_store_heterogeneous(pool_size)
-            for v in range(self._num_variants):
-                if not self._variant_layouts[v]:
+            for env_id, pool in self._layout_pools.items():
+                if not pool:
                     raise RuntimeError(
-                        f"Placement pool failed to produce any valid layouts for variant {v} "
+                        f"Placement pool failed to produce any valid layouts for env {env_id} "
                         f"from {pool_size} attempts. Check object relations and constraints."
                     )
         else:
@@ -97,25 +93,9 @@ def __init__(
 
     @property
     def is_heterogeneous(self) -> bool:
-        """Whether this pool operates in heterogeneous (per-variant) mode."""
+        """Whether this pool operates in heterogeneous (per-env) mode."""
         return self._heterogeneous
 
-    # ------------------------------------------------------------------
-    # Variant helpers
-    # ------------------------------------------------------------------
-
-    @staticmethod
-    def _detect_num_variants(objects: list[ObjectBase]) -> int:
-        """Return the number of unique object variants across heterogeneous objects."""
-        for obj in objects:
-            if getattr(obj, "heterogeneous_bbox", False):
-                return len(obj.objects)  # type: ignore[attr-defined]
-        return 1
-
-    def _variant_for_env(self, env_id: int) -> int:
-        """Map an environment index to its variant index."""
-        return env_id % self._num_variants
-
     # ------------------------------------------------------------------
     # Homogeneous (flat pool) internals
     # ------------------------------------------------------------------
@@ -165,48 +145,73 @@ def _refill_pool_via_solve_if_required(self, count: int) -> None:
             )
 
     # ------------------------------------------------------------------
-    # Heterogeneous (per-variant sub-pool) internals
+    # Heterogeneous (per-env pool) internals
     # ------------------------------------------------------------------
 
-    def _compact_variant(self, variant: int) -> None:
-        """Drop consumed layouts for a single variant and reset its read index."""
-        idx = self._variant_next_idx[variant]
-        self._variant_layouts[variant] = self._variant_layouts[variant][idx:]
-        self._variant_next_idx[variant] = 0
+    def _compact_env_pool(self, env_id: int) -> None:
+        """Drop consumed layouts for a single env and reset its cursor."""
+        idx = self._layout_cursors[env_id]
+        self._layout_pools[env_id] = self._layout_pools[env_id][idx:]
+        self._layout_cursors[env_id] = 0
 
     def _solve_and_store_heterogeneous(self, num_layouts: int) -> None:
-        """Solve layouts and bucket valid results by variant index.
+        """Solve layouts and store valid results into per-env pools.
 
-        Each result ``i`` from the solver corresponds to variant
-        ``i % num_variants`` because ``get_bounding_box_per_env`` assigns
-        variants in round-robin order.
+        Each round solves ``num_envs`` layouts in one batched call.
+        Result ``i`` is solved with env ``i``'s actual object geometry
+        (from ``get_bounding_box_per_env(num_envs)``) and is stored
+        directly into ``_layout_pools[i]``.  Multiple rounds are run
+        until each env has enough candidates.
         """
-        for v in range(self._num_variants):
-            self._compact_variant(v)
-
-        with torch.inference_mode(False):
-            result = self._placer.place(self._objects, num_envs=num_layouts, result_per_env=True)
-
-        all_results = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
+        for env_id in self._layout_pools:
+            self._compact_env_pool(env_id)
 
+        num_rounds = max(1, num_layouts // self._num_envs)
         total_valid = 0
-        for i, r in enumerate(all_results):
-            if r.success:
-                variant = i % self._num_variants
-                self._variant_layouts[variant].append(r)
-                total_valid += 1
-
-        if total_valid < num_layouts:
-            print(
-                f"Placement pool (heterogeneous): solved {num_layouts} candidates,"
-                f" {total_valid} valid, {num_layouts - total_valid} failed validation"
+        total_solved = 0
+        all_round_results: list[list[PlacementResult]] = []
+
+        for _ in range(num_rounds):
+            with torch.inference_mode(False):
+                result = self._placer.place(self._objects, num_envs=self._num_envs, result_per_env=True)
+
+            round_results = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
+            all_round_results.append(round_results)
+            total_solved += len(round_results)
+
+            for env_id, r in enumerate(round_results):
+                if r.success:
+                    self._layout_pools[env_id].append(r)
+                    total_valid += 1
+
+        if total_valid < total_solved:
+            failed_envs = [e for e in self._layout_pools if not self._layout_pools[e]]
+            msg = (
+                f"Placement pool (heterogeneous): solved {total_solved} candidates"
+                f" across {num_rounds} rounds,"
+                f" {total_valid} valid, {total_solved - total_valid} failed validation"
             )
-
-        if total_valid == 0:
-            print("Warning: No candidates passed strict validation. Accepting best-loss layouts as fallback.")
-            for i, r in enumerate(all_results):
-                variant = i % self._num_variants
-                self._variant_layouts[variant].append(r)
+            if failed_envs:
+                msg += f". Envs with zero valid layouts: {failed_envs}"
+            print(msg)
+
+        # Per-env fallback: for any env that still has zero valid layouts,
+        # accept the best-loss (lowest loss) result from all rounds.
+        for env_id in range(self._num_envs):
+            if self._layout_pools[env_id]:
+                continue
+            best: PlacementResult | None = None
+            for round_results in all_round_results:
+                if env_id < len(round_results):
+                    r = round_results[env_id]
+                    if best is None or r.final_loss < best.final_loss:
+                        best = r
+            if best is not None:
+                print(
+                    f"Warning: env {env_id} had no valid layouts; "
+                    f"accepting best-loss fallback (loss={best.final_loss:.6f})."
+                )
+                self._layout_pools[env_id].append(best)
 
     # ------------------------------------------------------------------
     # Public API
@@ -221,7 +226,7 @@ def sample_without_replacement(
         read index.
 
         In **heterogeneous mode** ``env_ids`` must be provided so each
-        environment receives a layout matching its variant geometry.
+        environment receives a layout matching its object geometry.
 
         Args:
             count: Number of layouts to return.
@@ -244,6 +249,7 @@ def sample_without_replacement(
     def _sample_without_replacement_heterogeneous(
         self, count: int, env_ids: list[int] | torch.Tensor | None
     ) -> list[PlacementResult]:
+        """Draw one layout per requested env, refilling any depleted per-env pool."""
         assert env_ids is not None, "env_ids must be provided for heterogeneous placement pools."
 
         if isinstance(env_ids, torch.Tensor):
@@ -252,28 +258,32 @@ def _sample_without_replacement_heterogeneous(
             ids = list(env_ids)
         assert len(ids) == count
 
-        variant_demand: dict[int, int] = {}
-        for eid in ids:
-            v = self._variant_for_env(eid)
-            variant_demand[v] = variant_demand.get(v, 0) + 1
+        # Refill any env pool that doesn't have enough layouts.
+        demand_per_env: dict[int, int] = {}
+        for env_id in ids:
+            demand_per_env[env_id] = demand_per_env.get(env_id, 0) + 1
+
+        needs_refill = False
+        for env_id, demand in demand_per_env.items():
+            available = len(self._layout_pools[env_id]) - self._layout_cursors[env_id]
+            if available < demand:
+                needs_refill = True
+                break
 
-        for v, demand in variant_demand.items():
-            avail = len(self._variant_layouts[v]) - self._variant_next_idx[v]
-            if avail < demand:
-                refill = max(self._pool_size, demand * self._num_variants)
-                self._solve_and_store_heterogeneous(refill)
+        if needs_refill:
+            max_demand = max(demand_per_env.values())
+            self._solve_and_store_heterogeneous(max(self._pool_size, max_demand * self._num_envs))
 
         results: list[PlacementResult] = []
-        for eid in ids:
-            v = self._variant_for_env(eid)
-            idx = self._variant_next_idx[v]
-            if idx >= len(self._variant_layouts[v]):
+        for env_id in ids:
+            idx = self._layout_cursors[env_id]
+            if idx >= len(self._layout_pools[env_id]):
                 raise RuntimeError(
-                    f"Placement pool: variant {v} has no more valid layouts "
-                    f"(needed for env {eid}). The solver is not producing enough valid placements."
+                    f"Placement pool: env {env_id} has no more valid layouts. "
+                    "The solver is not producing enough valid placements."
                 )
-            results.append(self._variant_layouts[v][idx])
-            self._variant_next_idx[v] = idx + 1
+            results.append(self._layout_pools[env_id][idx])
+            self._layout_cursors[env_id] = idx + 1
 
         return results
 
@@ -281,8 +291,7 @@ def sample_with_replacement(self, count: int) -> list[PlacementResult]:
         """Pick *count* layouts at random with replacement (non-consuming).
 
         In **heterogeneous mode**, each position ``i`` in the returned
-        list corresponds to env ``i`` and is drawn from the sub-pool
-        matching that env's variant (``i % num_variants``).
+        list corresponds to env ``i`` and is drawn from that env's pool.
 
         Used by ``resolve_on_reset=False`` to assign initial positions
         that persist across resets.
@@ -292,11 +301,11 @@ def sample_with_replacement(self, count: int) -> list[PlacementResult]:
         return random.choices(self._layouts, k=count)
 
     def _sample_with_replacement_heterogeneous(self, count: int) -> list[PlacementResult]:
+        """Pick one random layout per env from its pool (non-consuming)."""
         results: list[PlacementResult] = []
-        for env_idx in range(count):
-            v = self._variant_for_env(env_idx)
-            pool = self._variant_layouts[v]
-            assert pool, f"Variant {v} has no valid layouts to sample from."
+        for env_id in range(count):
+            pool = self._layout_pools[env_id]
+            assert pool, f"Env {env_id} has no valid layouts to sample from."
             results.append(random.choice(pool))
         return results
 
@@ -304,10 +313,10 @@ def _sample_with_replacement_heterogeneous(self, count: int) -> list[PlacementRe
     def remaining(self) -> int:
         """Number of layouts not yet consumed by :meth:`sample_without_replacement`.
 
-        For heterogeneous pools, returns the minimum across all variants.
+        For heterogeneous pools, returns the minimum across all envs.
         """
         if self._heterogeneous:
-            return min(len(self._variant_layouts[v]) - self._variant_next_idx[v] for v in range(self._num_variants))
+            return min(len(self._layout_pools[e]) - self._layout_cursors[e] for e in self._layout_pools)
         return len(self._layouts) - self._next_idx
 
     @property
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index ab49efc99..a809051ba 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -367,3 +367,122 @@ def test_pooled_placer_homogeneous_unaffected_by_num_envs():
     assert not pool.is_heterogeneous
     draws = pool.sample_without_replacement(3)
     assert len(draws) == 3
+
+
+# ---------------------------------------------------------------------------
+# Multi-set heterogeneous: different variant counts across objects
+# ---------------------------------------------------------------------------
+
+
+def test_pooled_placer_multi_set_different_variant_counts():
+    """Pool with two heterogeneous objects having different variant counts.
+
+    Bottles (3 variants) and boxes (2 variants) across 6 envs.
+    Each env gets its own pool with layouts matching its object geometry.
+    """
+    desk = _make_desk()
+
+    bottle_small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.2))
+    bottle_medium = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.25))
+    bottle_large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.3))
+    bottles = HeterogeneousDummyObject(name="bottles", bboxes=[bottle_small, bottle_medium, bottle_large])
+    bottles.add_relation(On(desk, clearance_m=0.01))
+
+    box_small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.1, 0.1))
+    box_large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.15, 0.15))
+    boxes = HeterogeneousDummyObject(name="boxes", bboxes=[box_small, box_large])
+    boxes.add_relation(On(desk, clearance_m=0.01))
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
+
+    pool = PooledObjectPlacer(objects=[desk, bottles, boxes], placer_params=placer_params, pool_size=50, num_envs=6)
+    assert pool.is_heterogeneous
+    assert pool.remaining > 0
+
+    env_ids = torch.tensor([0, 1, 2, 3, 4, 5])
+    draws = pool.sample_without_replacement(6, env_ids=env_ids)
+    assert len(draws) == 6
+    for d in draws:
+        assert bottles in d.positions
+        assert boxes in d.positions
+
+
+def test_pooled_placer_multi_set_sample_with_replacement():
+    """sample_with_replacement with multi-set heterogeneous objects."""
+    desk = _make_desk()
+
+    a_s = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.15))
+    a_m = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.2))
+    a_l = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.25))
+    obj_a = HeterogeneousDummyObject(name="A", bboxes=[a_s, a_m, a_l])
+    obj_a.add_relation(On(desk, clearance_m=0.01))
+
+    b_s = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
+    b_l = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.15, 0.12))
+    obj_b = HeterogeneousDummyObject(name="B", bboxes=[b_s, b_l])
+    obj_b.add_relation(On(desk, clearance_m=0.01))
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
+
+    pool = PooledObjectPlacer(objects=[desk, obj_a, obj_b], placer_params=placer_params, pool_size=50, num_envs=6)
+
+    initial_remaining = pool.remaining
+    samples = pool.sample_with_replacement(6)
+    assert len(samples) == 6
+    assert pool.remaining == initial_remaining
+
+
+def test_pooled_placer_multi_set_refill():
+    """Exhausting a per-env pool should trigger refill with multi-set objects."""
+    desk = _make_desk()
+
+    v1 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.15))
+    v2 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.12, 0.12, 0.2))
+    v3 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.18))
+    obj_a = HeterogeneousDummyObject(name="A", bboxes=[v1, v2, v3])
+    obj_a.add_relation(On(desk, clearance_m=0.01))
+
+    w1 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.1, 0.1))
+    w2 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.18, 0.12, 0.12))
+    obj_b = HeterogeneousDummyObject(name="B", bboxes=[w1, w2])
+    obj_b.add_relation(On(desk, clearance_m=0.01))
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
+
+    # 6 envs with 3+2 variant objects, small pool to force refill
+    pool = PooledObjectPlacer(objects=[desk, obj_a, obj_b], placer_params=placer_params, pool_size=12, num_envs=6)
+
+    # Drain pool then request more to trigger refill
+    initial_remaining = pool.remaining
+    env_ids = torch.tensor(list(range(6)) * initial_remaining)
+    pool.sample_without_replacement(len(env_ids), env_ids=env_ids)
+
+    env_ids_more = torch.tensor([0, 1, 2, 3, 4, 5])
+    draws = pool.sample_without_replacement(6, env_ids=env_ids_more)
+    assert len(draws) == 6
+
+
+def test_pooled_placer_per_env_pools_isolated():
+    """Each env_id should have its own independent pool of layouts."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
+
+    initial_remaining = pool.remaining
+
+    # Draw only from env 0 and env 1; env 2 and 3 should be unaffected.
+    env_ids = torch.tensor([0, 1])
+    pool.sample_without_replacement(2, env_ids=env_ids)
+
+    # `remaining` reports the min across all envs.  Env 0 and 1 each lost one
+    # layout, so the min should have decreased by 1.
+    assert pool.remaining == initial_remaining - 1
+
+    # Drawing from env 2 and 3 should still work from their full pools.
+    env_ids_23 = torch.tensor([2, 3])
+    draws = pool.sample_without_replacement(2, env_ids=env_ids_23)
+    assert len(draws) == 2
+    for d in draws:
+        assert hetero in d.positions

From b1bef41084ed2c1ffee4bf0f21d2c7d8e22a60cb Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 4 May 2026 12:55:10 -0700
Subject: [PATCH 06/46] fix falling objects

---
 isaaclab_arena/relations/relation_solver.py   |  33 ++--
 ...e_multi_object_no_collision_environment.py | 151 ++++++++++++++++--
 2 files changed, 151 insertions(+), 33 deletions(-)

diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index 844c83a4f..917f5b0c5 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -152,12 +152,15 @@ def _compute_no_overlap_loss(
         debug: bool = False,
         bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> torch.Tensor:
-        """Compute pairwise no-overlap loss for all non-anchor objects against all other objects.
+        """Compute pairwise no-overlap loss for non-anchor object pairs only.
 
-        Each unique pair is evaluated twice (once per direction):
-        - Non-anchor vs anchor: gradient flows to the non-anchor only.
-        - Non-anchor vs non-anchor: both objects receive gradient by computing
-          the loss in both directions with the other's position detached.
+        Anchor objects (e.g. the table) are excluded: the On relation already
+        controls Z placement relative to anchors, and adding a 3D clearance
+        envelope around the anchor would fight the On constraint in Z.
+
+        Each unique non-anchor pair is evaluated twice (once per direction) so
+        both objects receive gradient.  Uses xy_only=True because objects on the
+        same surface share Z height.
 
         Args:
             state: Current optimization state with object positions.
@@ -171,26 +174,14 @@ def _compute_no_overlap_loss(
         total_loss = torch.zeros(state.batch_size, device=device, dtype=torch.float32)
 
         non_anchor_objects = state.optimizable_objects
-        anchor_objects = list(state.anchor_objects)
 
         for i, child in enumerate(non_anchor_objects):
             child_pos = state.get_position(child)
             child_bbox = self._get_bbox(child, device, bboxes_per_row)
 
-            # Against all anchors
-            for anchor in anchor_objects:
-                anchor_world_bbox = anchor.get_world_bounding_box().to(device)
-                loss = self._no_collision_strategy.compute_loss(
-                    clearance_m=self.params.clearance_m,
-                    child_pos=child_pos,
-                    child_bbox=child_bbox,
-                    parent_world_bbox=anchor_world_bbox,
-                )
-                if debug:
-                    print(f"  [NoOverlap] {child.name} vs {anchor.name}: loss={loss.mean().item():.6f}")
-                total_loss = total_loss + loss
-
-            # Against other non-anchors (unique pairs, both directions)
+            # Against other non-anchors (unique pairs, both directions).
+            # Uses xy_only=True because objects on the same surface share Z height;
+            # 3D overlap would generate Z-gradient fighting the On constraint.
             for j in range(i + 1, len(non_anchor_objects)):
                 other = non_anchor_objects[j]
                 other_pos = state.get_position(other)
@@ -203,6 +194,7 @@ def _compute_no_overlap_loss(
                     child_pos=child_pos,
                     child_bbox=child_bbox,
                     parent_world_bbox=other_world_bbox,
+                    xy_only=True,
                 )
 
                 # Reverse: gradient flows to other (object j)
@@ -212,6 +204,7 @@ def _compute_no_overlap_loss(
                     child_pos=other_pos,
                     child_bbox=other_bbox,
                     parent_world_bbox=child_world_bbox,
+                    xy_only=True,
                 )
 
                 if debug:
diff --git a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
index 61d0e1a1b..fb9071d5e 100644
--- a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
+++ b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
@@ -4,14 +4,40 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """
-Table + multi-object no-overlap environment. Office table with objects placed via
-On(table) with built-in no-overlap (relation solver). Includes a robot (e.g. GR1).
+Table + multi-object no-overlap environment with a robot (e.g. GR1).
 No task -- suitable for policy_runner with zero_action or any policy.
 
+Supports two placement modes via ``--mode``:
+
+* **homogeneous** (default): each object is a regular Object — same in all envs.
+* **heterogeneous**: objects are wrapped in ``RigidObjectSet`` for per-env variance.
+
+Both modes use the office table by default. Use ``--objects`` to override object
+lists for controlled experiments.
+
 Example (--viz kit enables the Kit visualizer, --episode_length_s triggers periodic resets):
+
+  # Homogeneous (default) — YCB objects
   /isaac-sim/python.sh isaaclab_arena/evaluation/policy_runner.py --viz kit --policy_type zero_action --num_steps 500 \\
     --num_envs 16 --env_spacing 4.0 --enable_cameras \\
     gr1_table_multi_object_no_collision --embodiment gr1_joint --episode_length_s 4.0
+
+  # Heterogeneous — robolab objects in RigidObjectSet
+  /isaac-sim/python.sh isaaclab_arena/evaluation/policy_runner.py --viz kit --policy_type zero_action --num_steps 500 \\
+    --num_envs 16 --env_spacing 4.0 --enable_cameras \\
+    gr1_table_multi_object_no_collision --embodiment gr1_joint --episode_length_s 4.0 --mode heterogeneous
+
+  # Experiment: robolab objects in HOMOGENEOUS mode (isolate object vs mode)
+  /isaac-sim/python.sh isaaclab_arena/evaluation/policy_runner.py --viz kit --policy_type zero_action --num_steps 500 \\
+    --num_envs 16 --env_spacing 4.0 --enable_cameras \\
+    gr1_table_multi_object_no_collision --embodiment gr1_joint --episode_length_s 4.0 \\
+    --objects ketchup_bottle_hope_robolab mustard_bottle_hope_robolab popcorn_box_hope_robolab
+
+  # Experiment: YCB objects in HETEROGENEOUS mode (isolate wrapping vs objects)
+  /isaac-sim/python.sh isaaclab_arena/evaluation/policy_runner.py --viz kit --policy_type zero_action --num_steps 500 \\
+    --num_envs 16 --env_spacing 4.0 --enable_cameras \\
+    gr1_table_multi_object_no_collision --embodiment gr1_joint --episode_length_s 4.0 \\
+    --mode heterogeneous --objects cracker_box sugar_box tomato_soup_can mustard_bottle
 """
 
 from __future__ import annotations
@@ -37,12 +63,54 @@
 # objects. Better initialization strategies and constraining unchanged pose dimensions
 # are needed in the near future.
 
+# -- Heterogeneous mode default object sets ----------------------------------
+# Each entry is a multi-variant RigidObjectSet — each env gets a different
+# variant. Objects sourced from het-viz branch gif capture script.
+HETERO_VARIANT_SETS = {
+    "bottles": [
+        "mustard_bottle_hope_robolab",
+        "milk_carton_hope_robolab",
+        "orange_juice_carton_hope_robolab",
+        "parmesan_cheese_canister_hope_robolab",
+    ],
+    "cans": [
+        "alphabet_soup_can_hope_robolab",
+        "canned_peaches_hope_robolab",
+        "corn_can_hope_robolab",
+        "tomato_sauce_can_hope_robolab",
+        "pineapple_slices_can_hope_robolab",
+        "green_beans_can_hope_robolab",
+    ],
+    "tools": [
+        "spoon_handal_robolab",
+        "spoon_1_handal_robolab",
+        "spoon_2_handal_robolab",
+        "measuring_spoon_handal_robolab",
+    ],
+    "boxes": [
+        "popcorn_box_hope_robolab",
+        "chocolate_pudding_mix_hope_robolab",
+        "macaroni_and_cheese_hope_robolab",
+        "granola_bars_hope_robolab",
+    ],
+}
+
+HETERO_FIXED_OBJECTS = [
+    ("banana_ycb_robolab", 0.5, -0.15),
+    ("lime01_fruits_veggies_robolab", 0.5, 0.15),
+]
+
+
+
 
 @register_environment
 class GR1TableMultiObjectNoCollisionEnvironment(ExampleEnvironmentBase):
     """
     Table-based scene with multiple objects (On(table) + built-in no-overlap) and a robot.
     Layout is solved by ArenaEnvBuilder default relation solving; reset uses asset events.
+
+    Supports ``--mode homogeneous`` (default) and ``--mode heterogeneous`` for
+    inter-environment object variance via ``RigidObjectSet``.
     """
 
     name: str = "gr1_table_multi_object_no_collision"
@@ -50,11 +118,12 @@ class GR1TableMultiObjectNoCollisionEnvironment(ExampleEnvironmentBase):
     def get_env(self, args_cli: argparse.Namespace) -> IsaacLabArenaEnvironment:
         from isaaclab_arena.assets.object_reference import ObjectReference
         from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment
-        from isaaclab_arena.relations.relations import IsAnchor, On
+        from isaaclab_arena.relations.relations import IsAnchor
         from isaaclab_arena.scene.scene import Scene
         from isaaclab_arena.tasks.no_task import NoTask
         from isaaclab_arena.utils.pose import Pose
 
+        mode = getattr(args_cli, "mode", "homogeneous")
         enable_cameras = getattr(args_cli, "enable_cameras", False)
         camera_offset = Pose(
             position_xyz=(0.12515, 0.0, 0.06776),
@@ -73,10 +142,9 @@ def get_env(self, args_cli: argparse.Namespace) -> IsaacLabArenaEnvironment:
         )
 
         ground_plane = self.asset_registry.get_asset_by_name("ground_plane")()
-        table_background = self.asset_registry.get_asset_by_name("office_table")()
         light = self.asset_registry.get_asset_by_name("light")()
 
-        # Table surface as anchor for On relations
+        table_background = self.asset_registry.get_asset_by_name("office_table")()
         tabletop_reference = ObjectReference(
             name="table",
             prim_path="{ENV_REGEX_NS}/office_table/Geometry/sm_tabletop_a01_01/sm_tabletop_a01_top_01",
@@ -84,13 +152,11 @@ def get_env(self, args_cli: argparse.Namespace) -> IsaacLabArenaEnvironment:
         )
         tabletop_reference.add_relation(IsAnchor())
 
-        object_names = getattr(args_cli, "objects", None) or DEFAULT_TABLE_OBJECTS
-        placeable_assets = []
-        for name in object_names:
-            obj = self.asset_registry.get_asset_by_name(name)()
-            obj.add_relation(On(tabletop_reference))
-            placeable_assets.append(obj)
-        # No-overlap between all pairs is handled automatically by the solver (built-in behavior).
+        object_names = getattr(args_cli, "objects", None)
+        if mode == "heterogeneous":
+            placeable_assets = self._build_heterogeneous_objects(tabletop_reference, object_names)
+        else:
+            placeable_assets = self._build_homogeneous_objects(tabletop_reference, object_names)
 
         if args_cli.teleop_device is not None:
             teleop_device = self.device_registry.get_device_by_name(args_cli.teleop_device)()
@@ -131,6 +197,54 @@ def _enable_periodic_reset(cfg):
         )
         return isaaclab_arena_environment
 
+    def _build_homogeneous_objects(self, tabletop_reference, object_names=None):
+        """Build placeable objects for homogeneous mode (same objects in all envs).
+
+        Each object is a regular Object instance — identical across environments.
+        """
+        from isaaclab_arena.relations.relations import On
+
+        names = object_names or DEFAULT_TABLE_OBJECTS
+        placeable_assets = []
+        for name in names:
+            obj = self.asset_registry.get_asset_by_name(name)()
+            obj.add_relation(On(tabletop_reference, clearance_m=0.001))
+            placeable_assets.append(obj)
+        return placeable_assets
+
+    def _build_heterogeneous_objects(self, tabletop_reference, object_names=None):
+        """Build placeable objects for heterogeneous mode.
+
+        When --objects is provided, each object becomes a single-variant RigidObjectSet.
+        Otherwise, uses HETERO_FIXED_OBJECTS (pinned fruits) + HETERO_VARIANT_SETS
+        (multi-variant sets from het-viz branch).
+        """
+        from isaaclab_arena.assets.object_set import RigidObjectSet
+        from isaaclab_arena.relations.relations import AtPosition, On
+
+        if object_names:
+            placeable_assets = []
+            for name in object_names:
+                obj = self.asset_registry.get_asset_by_name(name)()
+                obj_set = RigidObjectSet(name=name, objects=[obj])
+                obj_set.add_relation(On(tabletop_reference, clearance_m=0.001))
+                placeable_assets.append(obj_set)
+        else:
+            placeable_assets = []
+            for name, x, y in HETERO_FIXED_OBJECTS:
+                obj = self.asset_registry.get_asset_by_name(name)()
+                obj.add_relation(On(tabletop_reference, clearance_m=0.001))
+                obj.add_relation(AtPosition(x=x, y=y))
+                placeable_assets.append(obj)
+
+            for set_name, variant_names in HETERO_VARIANT_SETS.items():
+                members = [self.asset_registry.get_asset_by_name(n)() for n in variant_names]
+                obj_set = RigidObjectSet(name=set_name, objects=members)
+                obj_set.add_relation(On(tabletop_reference, clearance_m=0.001))
+                placeable_assets.append(obj_set)
+
+        return placeable_assets
+
     @staticmethod
     def add_cli_args(parser: argparse.ArgumentParser) -> None:
         parser.add_argument(
@@ -138,7 +252,11 @@ def add_cli_args(parser: argparse.ArgumentParser) -> None:
             nargs="*",
             type=str,
             default=None,
-            help=f"Object names to spawn on the table (On + no-overlap). Default: {' '.join(DEFAULT_TABLE_OBJECTS)}",
+            help=(
+                "Object names (works in both modes). "
+                f"Homo default: {' '.join(DEFAULT_TABLE_OBJECTS)}; "
+                f"Hetero default: multi-variant sets (cracker_box/sugar_box, tomato_soup_can/mustard_bottle)"
+            ),
         )
         parser.add_argument("--embodiment", type=str, default="gr1_joint", help="Robot embodiment to use")
         parser.add_argument("--teleop_device", type=str, default=None, help="Teleoperation device to use")
@@ -148,3 +266,10 @@ def add_cli_args(parser: argparse.ArgumentParser) -> None:
             default=None,
             help="Episode length in seconds. Enables time_out termination so objects are re-placed on reset.",
         )
+        parser.add_argument(
+            "--mode",
+            type=str,
+            default="homogeneous",
+            choices=["homogeneous", "heterogeneous"],
+            help="Placement mode: 'homogeneous' (same objects everywhere) or 'heterogeneous' (per-env variants).",
+        )

From 18cf4457923ae328bb5bd2ec2f9d681b370b98c2 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 4 May 2026 16:39:22 -0700
Subject: [PATCH 07/46] clean up heter

---
 isaaclab_arena/relations/object_placer.py     | 25 +++++--
 .../relations/pooled_object_placer.py         | 66 +++++++++++--------
 2 files changed, 58 insertions(+), 33 deletions(-)

diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index 464c2a007..cf7c04fff 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -84,6 +84,7 @@ def place(
         objects: list[ObjectBase],
         num_envs: int = 1,
         result_per_env: bool = True,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> PlacementResult | MultiEnvPlacementResult:
         """Place objects according to their spatial relations.
 
@@ -95,6 +96,10 @@ def place(
             result_per_env: When True (default), each environment gets a distinct
                 layout. When False, a single best layout is solved and applied
                 identically to all environments.
+            env_bboxes: Pre-computed per-env bounding boxes (shape ``(num_envs, 3)``
+                per object).  When provided, ``_place_heterogeneous`` uses these
+                instead of calling ``get_bounding_box_per_env(num_envs)``.  This
+                allows the caller to tile real-env bboxes for pooled solving.
 
         Returns:
             PlacementResult when a single layout is produced (num_envs=1 or
@@ -138,7 +143,13 @@ def place(
 
         if heterogeneous:
             results_per_env = self._place_heterogeneous(
-                objects, anchor_objects_set, num_envs, max_attempts, num_candidates, generator
+                objects,
+                anchor_objects_set,
+                num_envs,
+                max_attempts,
+                num_candidates,
+                generator,
+                env_bboxes=env_bboxes,
             )
         else:
             results_per_env = self._place_homogeneous(
@@ -207,16 +218,20 @@ def _place_heterogeneous(
         max_attempts: int,
         num_candidates: int,
         generator: torch.Generator | None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> list[PlacementResult]:
         """Per-env placement: each candidate is tied to its env's object variants.
 
         Batch layout: candidates [e * max_attempts : (e+1) * max_attempts] belong
         to env *e*. Per-row bboxes reflect each env's actual variant geometry.
+
+        Args:
+            env_bboxes: When provided, uses these bboxes directly instead of
+                calling ``get_bounding_box_per_env(num_envs)``.  Each bbox must
+                have shape ``(num_envs, 3)``.
         """
-        # Build per-env bboxes (num_envs, 3) for every object.
-        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {
-            obj: obj.get_bounding_box_per_env(num_envs) for obj in objects
-        }
+        if env_bboxes is None:
+            env_bboxes = {obj: obj.get_bounding_box_per_env(num_envs) for obj in objects}
 
         # Expand into per-row bboxes (num_candidates, 3): repeat each env's
         # bbox max_attempts times so rows [e*A:(e+1)*A] share env e's geometry.
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 00507c03c..af36d879e 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -157,55 +157,65 @@ def _compact_env_pool(self, env_id: int) -> None:
     def _solve_and_store_heterogeneous(self, num_layouts: int) -> None:
         """Solve layouts and store valid results into per-env pools.
 
-        Each round solves ``num_envs`` layouts in one batched call.
-        Result ``i`` is solved with env ``i``'s actual object geometry
-        (from ``get_bounding_box_per_env(num_envs)``) and is stored
-        directly into ``_layout_pools[i]``.  Multiple rounds are run
-        until each env has enough candidates.
+        Computes bounding boxes for the real ``num_envs`` once, tiles them
+        to ``num_layouts`` entries, and solves everything in **one** batched
+        ``place()`` call.  Result ``i`` is mapped back to real env
+        ``i % num_envs`` for pool storage.
         """
+        from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
         for env_id in self._layout_pools:
             self._compact_env_pool(env_id)
 
-        num_rounds = max(1, num_layouts // self._num_envs)
-        total_valid = 0
-        total_solved = 0
-        all_round_results: list[list[PlacementResult]] = []
+        layouts_per_env = max(1, num_layouts // self._num_envs)
+        total_layouts = layouts_per_env * self._num_envs
 
-        for _ in range(num_rounds):
-            with torch.inference_mode(False):
-                result = self._placer.place(self._objects, num_envs=self._num_envs, result_per_env=True)
+        real_bboxes: dict = {obj: obj.get_bounding_box_per_env(self._num_envs) for obj in self._objects}
 
-            round_results = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
-            all_round_results.append(round_results)
-            total_solved += len(round_results)
+        tiled_bboxes: dict = {}
+        for obj, bbox in real_bboxes.items():
+            # (num_envs, 3) → repeat each env's row layouts_per_env times → (total_layouts, 3)
+            min_pt = bbox.min_point.repeat_interleave(layouts_per_env, dim=0)
+            max_pt = bbox.max_point.repeat_interleave(layouts_per_env, dim=0)
+            tiled_bboxes[obj] = AxisAlignedBoundingBox(min_point=min_pt, max_point=max_pt)
 
-            for env_id, r in enumerate(round_results):
-                if r.success:
-                    self._layout_pools[env_id].append(r)
-                    total_valid += 1
+        with torch.inference_mode(False):
+            result = self._placer.place(
+                self._objects,
+                num_envs=total_layouts,
+                result_per_env=True,
+                env_bboxes=tiled_bboxes,
+            )
+
+        all_results = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
+
+        total_valid = 0
+        for i, r in enumerate(all_results):
+            env_id = i // layouts_per_env
+            if r.success:
+                self._layout_pools[env_id].append(r)
+                total_valid += 1
 
+        total_solved = len(all_results)
         if total_valid < total_solved:
             failed_envs = [e for e in self._layout_pools if not self._layout_pools[e]]
             msg = (
-                f"Placement pool (heterogeneous): solved {total_solved} candidates"
-                f" across {num_rounds} rounds,"
+                f"Placement pool (heterogeneous): solved {total_solved} candidates,"
                 f" {total_valid} valid, {total_solved - total_valid} failed validation"
             )
             if failed_envs:
                 msg += f". Envs with zero valid layouts: {failed_envs}"
             print(msg)
 
-        # Per-env fallback: for any env that still has zero valid layouts,
-        # accept the best-loss (lowest loss) result from all rounds.
         for env_id in range(self._num_envs):
             if self._layout_pools[env_id]:
                 continue
             best: PlacementResult | None = None
-            for round_results in all_round_results:
-                if env_id < len(round_results):
-                    r = round_results[env_id]
-                    if best is None or r.final_loss < best.final_loss:
-                        best = r
+            start = env_id * layouts_per_env
+            end = start + layouts_per_env
+            for r in all_results[start:end]:
+                if best is None or r.final_loss < best.final_loss:
+                    best = r
             if best is not None:
                 print(
                     f"Warning: env {env_id} had no valid layouts; "

From 69579607b8f90f9450ee4ba778beb219edab02e0 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 15:11:08 -0700
Subject: [PATCH 08/46] rename env_bbox

---
 isaaclab_arena/relations/object_placer.py   |  8 ++---
 isaaclab_arena/relations/relation_solver.py | 34 ++++++++++-----------
 2 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index cf7c04fff..20d23c849 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -233,14 +233,14 @@ def _place_heterogeneous(
         if env_bboxes is None:
             env_bboxes = {obj: obj.get_bounding_box_per_env(num_envs) for obj in objects}
 
-        # Expand into per-row bboxes (num_candidates, 3): repeat each env's
+        # Expand into per-candidate bboxes (num_candidates, 3): repeat each env's
         # bbox max_attempts times so rows [e*A:(e+1)*A] share env e's geometry.
-        bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] = {}
+        candidate_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {}
         for obj, bbox in env_bboxes.items():
             # bbox.min_point is (num_envs, 3) → repeat_interleave → (num_candidates, 3)
             min_pt = bbox.min_point.repeat_interleave(max_attempts, dim=0)
             max_pt = bbox.max_point.repeat_interleave(max_attempts, dim=0)
-            bboxes_per_row[obj] = AxisAlignedBoundingBox(min_point=min_pt, max_point=max_pt)
+            candidate_bboxes[obj] = AxisAlignedBoundingBox(min_point=min_pt, max_point=max_pt)
 
         # Generate initial positions; each candidate uses its env's bbox.
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
@@ -260,7 +260,7 @@ def _place_heterogeneous(
                 self._generate_initial_positions(objects, anchor_objects_set, generator, child_bboxes=env_child_bboxes)
             )
 
-        all_positions = self._solver.solve(objects, initial_positions, bboxes_per_row=bboxes_per_row)
+        all_positions = self._solver.solve(objects, initial_positions, env_bboxes=candidate_bboxes)
         assert self._solver.last_loss_per_env is not None
         all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
 
diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index 917f5b0c5..4f2620152 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -71,25 +71,25 @@ def _get_bbox(
         self,
         obj: ObjectBase,
         device: torch.device | None,
-        bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] | None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None,
     ) -> AxisAlignedBoundingBox:
-        """Return the per-row or default local bbox for *obj*, moved to *device*."""
-        if bboxes_per_row is not None and obj in bboxes_per_row:
-            return bboxes_per_row[obj].to(device)
+        """Return the per-env or default local bbox for *obj*, moved to *device*."""
+        if env_bboxes is not None and obj in env_bboxes:
+            return env_bboxes[obj].to(device)
         return obj.get_bounding_box().to(device)
 
     def _compute_total_loss(
         self,
         state: RelationSolverState,
         debug: bool = False,
-        bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> torch.Tensor:
         """Compute total loss from all relations using registered strategies.
 
         Args:
             state: Current optimization state with object positions.
             debug: If True, print detailed loss breakdown.
-            bboxes_per_row: Optional per-row bboxes keyed by object. When
+            env_bboxes: Optional per-env bboxes keyed by object. When
                 provided the bbox ``min_point`` / ``max_point`` have shape
                 ``(batch, 3)`` instead of ``(1, 3)``, enabling heterogeneous
                 object placement.
@@ -106,7 +106,7 @@ def _compute_total_loss(
             for relation in obj.get_spatial_relations():
                 child_pos = state.get_position(obj)
                 strategy = self._get_strategy(relation)
-                child_bbox = self._get_bbox(obj, device, bboxes_per_row)
+                child_bbox = self._get_bbox(obj, device, env_bboxes)
 
                 # Handle unary relations (no parent)
                 if isinstance(relation, UnaryRelation):
@@ -124,7 +124,7 @@ def _compute_total_loss(
                         parent_world_bbox = parent.get_world_bounding_box().to(device)
                     else:
                         parent_pos = state.get_position(parent)
-                        parent_bbox = self._get_bbox(parent, device, bboxes_per_row)
+                        parent_bbox = self._get_bbox(parent, device, env_bboxes)
                         parent_world_bbox = parent_bbox.translated(parent_pos)
                     loss = strategy.compute_loss(
                         relation=relation,
@@ -141,7 +141,7 @@ def _compute_total_loss(
                 total_loss = total_loss + loss
 
         # Add built-in no-overlap loss between all object pairs
-        total_loss = total_loss + self._compute_no_overlap_loss(state, debug, bboxes_per_row=bboxes_per_row)
+        total_loss = total_loss + self._compute_no_overlap_loss(state, debug, env_bboxes=env_bboxes)
 
         self._last_loss_per_env = total_loss.detach().clone()
         return total_loss.mean()
@@ -150,7 +150,7 @@ def _compute_no_overlap_loss(
         self,
         state: RelationSolverState,
         debug: bool = False,
-        bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> torch.Tensor:
         """Compute pairwise no-overlap loss for non-anchor object pairs only.
 
@@ -165,7 +165,7 @@ def _compute_no_overlap_loss(
         Args:
             state: Current optimization state with object positions.
             debug: If True, print detailed loss breakdown.
-            bboxes_per_row: Optional per-row bboxes for heterogeneous placement.
+            env_bboxes: Optional per-env bboxes for heterogeneous placement.
 
         Returns:
             Per-environment loss tensor of shape (batch_size,).
@@ -177,7 +177,7 @@ def _compute_no_overlap_loss(
 
         for i, child in enumerate(non_anchor_objects):
             child_pos = state.get_position(child)
-            child_bbox = self._get_bbox(child, device, bboxes_per_row)
+            child_bbox = self._get_bbox(child, device, env_bboxes)
 
             # Against other non-anchors (unique pairs, both directions).
             # Uses xy_only=True because objects on the same surface share Z height;
@@ -185,7 +185,7 @@ def _compute_no_overlap_loss(
             for j in range(i + 1, len(non_anchor_objects)):
                 other = non_anchor_objects[j]
                 other_pos = state.get_position(other)
-                other_bbox = self._get_bbox(other, device, bboxes_per_row)
+                other_bbox = self._get_bbox(other, device, env_bboxes)
 
                 # Forward: gradient flows to child (object i)
                 other_world_bbox = other_bbox.translated(other_pos.detach())
@@ -220,7 +220,7 @@ def solve(
         self,
         objects: list[ObjectBase],
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]],
-        bboxes_per_row: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> list[dict[ObjectBase, tuple[float, float, float]]]:
         """Solve for optimal positions of all objects.
 
@@ -229,7 +229,7 @@ def solve(
                 marked with IsAnchor() which serves as a fixed reference.
             initial_positions: List of dicts (one per env). Use a single-element list
                 for single-env placement.
-            bboxes_per_row: Optional per-row bounding boxes keyed by object.
+            env_bboxes: Optional per-env bounding boxes keyed by object.
                 When provided, each ``AxisAlignedBoundingBox`` has shape
                 ``(batch, 3)`` so different batch rows can use different
                 geometry (heterogeneous placement). If ``None``, every row
@@ -263,7 +263,7 @@ def solve(
         # Compute initial loss so _last_loss_per_env is always populated
         # (needed even when max_iters=0, e.g. tests that only check init positions).
         with torch.no_grad():
-            self._compute_total_loss(state, bboxes_per_row=bboxes_per_row)
+            self._compute_total_loss(state, env_bboxes=env_bboxes)
 
         # Optimization loop
         loss_history = []
@@ -276,7 +276,7 @@ def solve(
                 position_history.append(state.get_all_positions_snapshot())
 
             # Compute total loss
-            loss = self._compute_total_loss(state, bboxes_per_row=bboxes_per_row)
+            loss = self._compute_total_loss(state, env_bboxes=env_bboxes)
             loss_history.append(loss.item())
 
             # Backprop and update (only optimizable positions will update)

From 78a511f69be296414bf8130100ecb57e9419d1cc Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 15:11:29 -0700
Subject: [PATCH 09/46] rename env_bbox

---
 isaaclab_arena/tests/test_heterogeneous_placement.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index a809051ba..69cbfdea7 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -88,12 +88,12 @@ def test_heterogeneous_dummy_returns_different_bboxes():
 
 
 # ---------------------------------------------------------------------------
-# Solver with per-row bboxes
+# Solver with per-env bboxes
 # ---------------------------------------------------------------------------
 
 
-def test_solver_accepts_per_row_bboxes():
-    """Solver should accept bboxes_per_row and produce valid results."""
+def test_solver_accepts_env_bboxes():
+    """Solver should accept env_bboxes and produce valid results."""
 
     desk = _make_desk()
     box = DummyObject(
@@ -107,14 +107,14 @@ def test_solver_accepts_per_row_bboxes():
 
     initial_positions = [{desk: (0.0, 0.0, 0.0), box: (0.5, 0.5, 0.11)} for _ in range(batch_size)]
 
-    # Create per-row bboxes with varying sizes across the batch.
+    # Create per-env bboxes with varying sizes across the batch.
     min_pts = torch.zeros(batch_size, 3)
     max_pts = torch.stack([torch.tensor([0.1 + 0.05 * i, 0.1 + 0.05 * i, 0.2]) for i in range(batch_size)])
-    per_row_bbox = AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts)
+    env_bbox = AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts)
 
     solver_params = RelationSolverParams(max_iters=100, convergence_threshold=1e-3, verbose=False)
     solver = RelationSolver(params=solver_params)
-    result = solver.solve(objects, initial_positions, bboxes_per_row={box: per_row_bbox})
+    result = solver.solve(objects, initial_positions, env_bboxes={box: env_bbox})
 
     assert len(result) == batch_size
     for pos_dict in result:

From 55cb2c9c38851318c09f7ca8ccdf2027465261af Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 15:37:43 -0700
Subject: [PATCH 10/46] Add force_convex_hull and xy_only

---
 .../environments/arena_env_builder.py         | 22 ++++++++++++
 .../relations/relation_loss_strategies.py     | 34 ++++++++++++-------
 ...e_multi_object_no_collision_environment.py |  4 +--
 3 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py
index f0f4a5b8f..c4c5f44ff 100644
--- a/isaaclab_arena/environments/arena_env_builder.py
+++ b/isaaclab_arena/environments/arena_env_builder.py
@@ -292,7 +292,29 @@ def make_registered_and_return_cfg(
     ) -> tuple[ManagerBasedEnv, IsaacLabArenaManagerBasedRLEnvCfg]:
         name, cfg = self.build_registered(env_cfg)
         env = gym.make(name, cfg=cfg, render_mode=render_mode)
+        _force_convex_hull(env)
         # ViewportCameraController sets the camera before KitVisualizer.initialize() is called,
         # so the call is silently ignored. Re-apply here once the visualizers are fully initialized.
         reapply_viewer_cfg(env)
         return env, cfg
+
+
+def _force_convex_hull(env: ManagerBasedEnv) -> None:
+    """Replace ``convexDecomposition`` with ``convexHull`` on all MeshCollision prims.
+
+    ``convexDecomposition`` on raw scanned meshes (e.g. robolab assets) creates
+    irregular contact surfaces that are unstable in multi-object scenarios.
+    ``convexHull`` produces a single convex shape that behaves predictably.
+    """
+    from pxr import UsdPhysics
+
+    stage = env.unwrapped.sim.stage
+    for prim in stage.Traverse():
+        if not prim.HasAPI(UsdPhysics.MeshCollisionAPI):
+            continue
+        mesh_col = UsdPhysics.MeshCollisionAPI(prim)
+        approx_attr = mesh_col.GetApproximationAttr()
+        if not approx_attr or not approx_attr.HasValue():
+            continue
+        if approx_attr.Get() == "convexDecomposition":
+            approx_attr.Set("convexHull")
diff --git a/isaaclab_arena/relations/relation_loss_strategies.py b/isaaclab_arena/relations/relation_loss_strategies.py
index af42a2e98..973854037 100644
--- a/isaaclab_arena/relations/relation_loss_strategies.py
+++ b/isaaclab_arena/relations/relation_loss_strategies.py
@@ -348,6 +348,7 @@ def compute_loss(
         child_pos: torch.Tensor,
         child_bbox: AxisAlignedBoundingBox,
         parent_world_bbox: AxisAlignedBoundingBox,
+        xy_only: bool = False,
     ) -> torch.Tensor:
         """Compute loss for no-overlap constraint.
 
@@ -356,6 +357,8 @@ def compute_loss(
             child_pos: Child object position (N, 3) in world coords.
             child_bbox: Child object local bounding box (N=1).
             parent_world_bbox: Parent bounding box in world coordinates.
+            xy_only: If True, compute 2D (XY) overlap area instead of 3D volume.
+                Use for objects on the same surface where Z overlap is expected.
 
         Returns:
             Loss tensor of shape (N,).
@@ -370,8 +373,6 @@ def compute_loss(
         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
 
         # Child world extents
         child_world_min = child_pos + child_bbox.min_point
@@ -380,11 +381,18 @@ def compute_loss(
         # 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)
         overlap_y = interval_overlap_axis_loss(child_world_min[:, 1], child_world_max[:, 1], parent_y_min, parent_y_max)
-        overlap_z = interval_overlap_axis_loss(child_world_min[:, 2], child_world_max[:, 2], parent_z_min, parent_z_max)
 
-        # 2. Volume loss: slope * product of per-axis overlap lengths (overlap volume when slope 1.0)
-        overlap_volume = overlap_x * overlap_y * overlap_z
-        total_loss = self.slope * overlap_volume
+        if xy_only:
+            overlap_area = overlap_x * overlap_y
+            total_loss = self.slope * overlap_area
+        else:
+            parent_z_min = parent_world_bbox.min_point[:, 2] - c
+            parent_z_max = parent_world_bbox.max_point[:, 2] + c
+            overlap_z = interval_overlap_axis_loss(
+                child_world_min[:, 2], child_world_max[:, 2], parent_z_min, parent_z_max
+            )
+            overlap_volume = overlap_x * overlap_y * overlap_z
+            total_loss = self.slope * overlap_volume
 
         if self.debug and child_pos.shape[0] == 1:
             print(
@@ -397,12 +405,14 @@ def compute_loss(
                 f" {child_world_max[0, 1].item():.4f}], parent_y=[{parent_y_min[0].item():.4f},"
                 f" {parent_y_max[0].item():.4f}])"
             )
-            print(
-                f"    [NoCollision] Z: overlap={overlap_z[0].item():.6f} (child_z=[{child_world_min[0, 2].item():.4f},"
-                f" {child_world_max[0, 2].item():.4f}], parent_z=[{parent_z_min[0].item():.4f},"
-                f" {parent_z_max[0].item():.4f}])"
-            )
-            print(f"    [NoCollision] volume={overlap_volume[0].item():.6f}, loss={total_loss[0].item():.6f}")
+            if not xy_only:
+                print(
+                    f"    [NoCollision] Z: overlap={overlap_z[0].item():.6f}"
+                    f" (child_z=[{child_world_min[0, 2].item():.4f},"
+                    f" {child_world_max[0, 2].item():.4f}], parent_z=[{parent_z_min[0].item():.4f},"
+                    f" {parent_z_max[0].item():.4f}])"
+                )
+            print(f"    [NoCollision] loss={total_loss[0].item():.6f}")
 
         return total_loss.squeeze(0) if single_input else total_loss
 
diff --git a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
index fb9071d5e..0046824ec 100644
--- a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
+++ b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
@@ -101,8 +101,6 @@
 ]
 
 
-
-
 @register_environment
 class GR1TableMultiObjectNoCollisionEnvironment(ExampleEnvironmentBase):
     """
@@ -255,7 +253,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> None:
             help=(
                 "Object names (works in both modes). "
                 f"Homo default: {' '.join(DEFAULT_TABLE_OBJECTS)}; "
-                f"Hetero default: multi-variant sets (cracker_box/sugar_box, tomato_soup_can/mustard_bottle)"
+                "Hetero default: multi-variant sets (cracker_box/sugar_box, tomato_soup_can/mustard_bottle)"
             ),
         )
         parser.add_argument("--embodiment", type=str, default="gr1_joint", help="Robot embodiment to use")

From 379b96b1d1dbc80bc14c941aee4a3ee8fdb58cdc Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 15:49:18 -0700
Subject: [PATCH 11/46] clean up comments

---
 ...r1_table_multi_object_no_collision_environment.py | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
index 0046824ec..280e9b909 100644
--- a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
+++ b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
@@ -26,18 +26,6 @@
   /isaac-sim/python.sh isaaclab_arena/evaluation/policy_runner.py --viz kit --policy_type zero_action --num_steps 500 \\
     --num_envs 16 --env_spacing 4.0 --enable_cameras \\
     gr1_table_multi_object_no_collision --embodiment gr1_joint --episode_length_s 4.0 --mode heterogeneous
-
-  # Experiment: robolab objects in HOMOGENEOUS mode (isolate object vs mode)
-  /isaac-sim/python.sh isaaclab_arena/evaluation/policy_runner.py --viz kit --policy_type zero_action --num_steps 500 \\
-    --num_envs 16 --env_spacing 4.0 --enable_cameras \\
-    gr1_table_multi_object_no_collision --embodiment gr1_joint --episode_length_s 4.0 \\
-    --objects ketchup_bottle_hope_robolab mustard_bottle_hope_robolab popcorn_box_hope_robolab
-
-  # Experiment: YCB objects in HETEROGENEOUS mode (isolate wrapping vs objects)
-  /isaac-sim/python.sh isaaclab_arena/evaluation/policy_runner.py --viz kit --policy_type zero_action --num_steps 500 \\
-    --num_envs 16 --env_spacing 4.0 --enable_cameras \\
-    gr1_table_multi_object_no_collision --embodiment gr1_joint --episode_length_s 4.0 \\
-    --mode heterogeneous --objects cracker_box sugar_box tomato_soup_can mustard_bottle
 """
 
 from __future__ import annotations

From fdb320346a0c35d0a962aec73223817e2c0416c0 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 16:01:24 -0700
Subject: [PATCH 12/46] address feedback

---
 isaaclab_arena/assets/object_set.py                  | 11 ++++++-----
 isaaclab_arena/relations/relation_loss_strategies.py |  7 +++++--
 isaaclab_arena/tests/test_heterogeneous_placement.py |  7 +++----
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 5cc1e1643..10e05d61e 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -71,11 +71,12 @@ def __init__(
         self.random_choice = random_choice
         self.heterogeneous_bbox: bool = len(objects) > 1
 
-        assert not (self.heterogeneous_bbox and self.random_choice), (
-            f"RigidObjectSet '{name}': random_choice=True is not supported with heterogeneous "
-            "placement (len(objects) > 1). The placement pool assumes round-robin variant "
-            "assignment (env_idx % num_variants) which conflicts with random spawning order."
-        )
+        if self.heterogeneous_bbox and self.random_choice:
+            raise ValueError(
+                f"RigidObjectSet '{name}': random_choice=True is not supported with heterogeneous "
+                "placement (len(objects) > 1). The placement pool assumes round-robin variant "
+                "assignment (env_idx % num_variants) which conflicts with random spawning order."
+            )
 
         # Set default prim_path if not provided
         if prim_path is None:
diff --git a/isaaclab_arena/relations/relation_loss_strategies.py b/isaaclab_arena/relations/relation_loss_strategies.py
index 973854037..b348e3c6f 100644
--- a/isaaclab_arena/relations/relation_loss_strategies.py
+++ b/isaaclab_arena/relations/relation_loss_strategies.py
@@ -325,8 +325,11 @@ class NoCollisionLossStrategy:
     Computes loss based on:
     1. X overlap: zero when child and parent are separated along X; else overlap length
     2. Y overlap: zero when separated along Y; else overlap length
-    3. Z overlap: zero when separated along Z; else overlap length
-    4. Volume loss: slope * (overlap_x * overlap_y * overlap_z)
+    3. Z overlap: zero when separated along Z; else overlap length (skipped when xy_only=True)
+    4. Loss: slope * overlap product (area when xy_only, volume otherwise)
+
+    When ``xy_only=True``, only XY overlap is used — suitable for objects on the
+    same surface where Z overlap is expected and Z gradients would fight the On constraint.
 
     This is a standalone strategy (not a RelationLossStrategy) because no-overlap
     is a built-in solver behavior, not a user-specified relation.
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 69cbfdea7..8c39ae3be 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -7,6 +7,8 @@
 
 import torch
 
+import pytest
+
 from isaaclab_arena.assets.dummy_object import DummyObject
 from isaaclab_arena.relations.object_placer import ObjectPlacer
 from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
@@ -316,11 +318,8 @@ def test_pooled_placer_heterogeneous_sample_without_replacement_requires_env_ids
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
-    try:
+    with pytest.raises(AssertionError):
         pool.sample_without_replacement(2, env_ids=None)
-        assert False, "Should have raised AssertionError"
-    except AssertionError:
-        pass
 
 
 def test_pooled_placer_heterogeneous_sample_with_replacement():

From 8482a9a0b795ad18a6a73e5a44506f69de850144 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 16:18:23 -0700
Subject: [PATCH 13/46] address feedback

---
 isaaclab_arena/environments/arena_env_builder.py          | 3 ++-
 isaaclab_arena/environments/isaaclab_arena_environment.py | 6 ++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py
index c4c5f44ff..bb74cc8cd 100644
--- a/isaaclab_arena/environments/arena_env_builder.py
+++ b/isaaclab_arena/environments/arena_env_builder.py
@@ -292,7 +292,8 @@ def make_registered_and_return_cfg(
     ) -> tuple[ManagerBasedEnv, IsaacLabArenaManagerBasedRLEnvCfg]:
         name, cfg = self.build_registered(env_cfg)
         env = gym.make(name, cfg=cfg, render_mode=render_mode)
-        _force_convex_hull(env)
+        if self.arena_env.force_convex_hull:
+            _force_convex_hull(env)
         # ViewportCameraController sets the camera before KitVisualizer.initialize() is called,
         # so the call is silently ignored. Re-apply here once the visualizers are fully initialized.
         reapply_viewer_cfg(env)
diff --git a/isaaclab_arena/environments/isaaclab_arena_environment.py b/isaaclab_arena/environments/isaaclab_arena_environment.py
index c8f2a9769..8498d5a56 100644
--- a/isaaclab_arena/environments/isaaclab_arena_environment.py
+++ b/isaaclab_arena/environments/isaaclab_arena_environment.py
@@ -29,6 +29,7 @@ def __init__(
         env_cfg_callback: Callable[IsaacLabArenaManagerBasedRLEnvCfg] | None = None,
         rl_framework_entry_point: str | None = None,
         rl_policy_cfg: str | None = None,
+        force_convex_hull: bool = False,
     ):
         """
         Args:
@@ -46,6 +47,10 @@ def __init__(
                 ``rl_policy_cfg`` is set.
             rl_policy_cfg: Import path to the RL policy config class, e.g.
                 ``"my_module:RLPolicyCfg"``.
+            force_convex_hull: If True, replace ``convexDecomposition`` with ``convexHull``
+                on all MeshCollision prims after scene creation. Needed for assets with
+                raw scanned meshes (e.g. robolab objects) that are unstable with
+                ``convexDecomposition``.
         """
         self.name = name
         self.scene = scene
@@ -53,6 +58,7 @@ def __init__(
         self.task = task
         self.teleop_device = teleop_device
         self.env_cfg_callback = env_cfg_callback
+        self.force_convex_hull = force_convex_hull
         if (rl_framework_entry_point is None) != (rl_policy_cfg is None):
             raise ValueError("rl_framework_entry_point and rl_policy_cfg must both be set or both be None.")
         self.rl_framework_entry_point = rl_framework_entry_point

From 8a7150f0b6443df0aaf528d63e9a77572a084c4b Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 18:11:40 -0700
Subject: [PATCH 14/46] test import

---
 .../tests/test_heterogeneous_placement.py     | 163 ++++++++++++++----
 ...e_multi_object_no_collision_environment.py |   7 +-
 2 files changed, 138 insertions(+), 32 deletions(-)

diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 8c39ae3be..4ba0f9f79 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -9,44 +9,48 @@
 
 import pytest
 
-from isaaclab_arena.assets.dummy_object import DummyObject
-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.pooled_object_placer import PooledObjectPlacer
-from isaaclab_arena.relations.relation_solver import RelationSolver
-from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
-from isaaclab_arena.relations.relations import IsAnchor, On
-from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-from isaaclab_arena.utils.pose import Pose
-
 # ---------------------------------------------------------------------------
-# Helpers
+# Helpers — all isaaclab_arena imports are deferred to avoid importing
+# Isaac Sim modules at pytest collection time.
 # ---------------------------------------------------------------------------
 
 
-class HeterogeneousDummyObject(DummyObject):
-    """DummyObject that provides different bounding boxes per environment.
+def _make_heterogeneous_dummy_class():
+    """Return the HeterogeneousDummyObject class (deferred import of DummyObject)."""
 
-    Used to exercise the heterogeneous placement path without requiring
-    RigidObjectSet's USD machinery.
-    """
+    from isaaclab_arena.assets.dummy_object import DummyObject
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+    class HeterogeneousDummyObject(DummyObject):
+        """DummyObject that provides different bounding boxes per environment.
+
+        Used to exercise the heterogeneous placement path without requiring
+        RigidObjectSet's USD machinery.
+        """
+
+        def __init__(self, name: str, bboxes: list[AxisAlignedBoundingBox], **kwargs):
+            super().__init__(name=name, bounding_box=bboxes[0], **kwargs)
+            self._per_env_bboxes = bboxes
+            self.heterogeneous_bbox = True
+            self.objects = bboxes
 
-    def __init__(self, name: str, bboxes: list[AxisAlignedBoundingBox], **kwargs):
-        super().__init__(name=name, bounding_box=bboxes[0], **kwargs)
-        self._per_env_bboxes = bboxes
-        self.heterogeneous_bbox = True
-        self.objects = bboxes
+        def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
+            n_variants = len(self._per_env_bboxes)
+            indices = [i % n_variants for i in range(num_envs)]
+            min_pts = torch.stack([self._per_env_bboxes[idx].min_point[0] for idx in indices])
+            max_pts = torch.stack([self._per_env_bboxes[idx].max_point[0] for idx in indices])
+            return AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts)
 
-    def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
-        n_variants = len(self._per_env_bboxes)
-        indices = [i % n_variants for i in range(num_envs)]
-        min_pts = torch.stack([self._per_env_bboxes[idx].min_point[0] for idx in indices])
-        max_pts = torch.stack([self._per_env_bboxes[idx].max_point[0] for idx in indices])
-        return AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts)
+    return HeterogeneousDummyObject
 
 
-def _make_desk() -> DummyObject:
+def _make_desk():
+    """Create a desk anchor object for tests."""
+    from isaaclab_arena.assets.dummy_object import DummyObject
+    from isaaclab_arena.relations.relations import IsAnchor
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+    from isaaclab_arena.utils.pose import Pose
+
     desk = DummyObject(
         name="desk",
         bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)),
@@ -64,6 +68,9 @@ def _make_desk() -> DummyObject:
 def test_dummy_object_bbox_per_env_expands_single():
     """Default get_bounding_box_per_env should repeat the single bbox."""
 
+    from isaaclab_arena.assets.dummy_object import DummyObject
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
     obj = DummyObject(
         name="box",
         bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)),
@@ -78,6 +85,10 @@ def test_dummy_object_bbox_per_env_expands_single():
 def test_heterogeneous_dummy_returns_different_bboxes():
     """HeterogeneousDummyObject should cycle through its member bboxes."""
 
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
+
     small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
     large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3))
     obj = HeterogeneousDummyObject(name="set", bboxes=[small, large])
@@ -97,6 +108,12 @@ def test_heterogeneous_dummy_returns_different_bboxes():
 def test_solver_accepts_env_bboxes():
     """Solver should accept env_bboxes and produce valid results."""
 
+    from isaaclab_arena.assets.dummy_object import DummyObject
+    from isaaclab_arena.relations.relation_solver import RelationSolver
+    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
+    from isaaclab_arena.relations.relations import On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
     desk = _make_desk()
     box = DummyObject(
         name="box",
@@ -132,6 +149,15 @@ def test_solver_accepts_env_bboxes():
 def test_placer_heterogeneous_produces_per_env_results():
     """Placer should detect heterogeneous objects and solve per-env."""
 
+    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.relation_solver_params import RelationSolverParams
+    from isaaclab_arena.relations.relations import On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
+
     desk = _make_desk()
 
     small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
@@ -161,6 +187,15 @@ def test_placer_heterogeneous_produces_per_env_results():
 def test_placer_heterogeneous_z_height_matches_variant():
     """Objects should be placed at z-height matching their env's variant bbox."""
 
+    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.relation_solver_params import RelationSolverParams
+    from isaaclab_arena.relations.relations import On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
+
     desk = _make_desk()
 
     # "tall" variant: height 0.4 → bottom at z ≈ 0.11 (desk top 0.1 + clearance 0.01)
@@ -199,6 +234,16 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
     bboxes per env while X is identical everywhere.
     """
 
+    from isaaclab_arena.assets.dummy_object import DummyObject
+    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.relation_solver_params import RelationSolverParams
+    from isaaclab_arena.relations.relations import On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
+
     desk = _make_desk()
 
     # A: heterogeneous — small variant in even envs, large in odd envs.
@@ -214,8 +259,6 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
     )
     obj_x.add_relation(On(desk, clearance_m=0.01))
 
-    # No-overlap is handled automatically by the solver's built-in clearance.
-
     objects = [desk, obj_a, obj_x]
     num_envs = 4
 
@@ -253,6 +296,14 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
 def test_homogeneous_path_unchanged():
     """When no heterogeneous objects exist, the homogeneous path is used."""
 
+    from isaaclab_arena.assets.dummy_object import DummyObject
+    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.relation_solver_params import RelationSolverParams
+    from isaaclab_arena.relations.relations import On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
     desk = _make_desk()
     box = DummyObject(
         name="box",
@@ -281,6 +332,13 @@ def test_homogeneous_path_unchanged():
 
 def _make_hetero_pool_objects():
     """Create desk + heterogeneous box for pool tests."""
+    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
+    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
+    from isaaclab_arena.relations.relations import On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
+
     desk = _make_desk()
     small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
     large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3))
@@ -294,6 +352,8 @@ def _make_hetero_pool_objects():
 
 def test_pooled_placer_heterogeneous_is_detected():
     """PooledObjectPlacer should detect heterogeneous objects and create variant sub-pools."""
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
@@ -303,6 +363,8 @@ def test_pooled_placer_heterogeneous_is_detected():
 
 def test_pooled_placer_heterogeneous_sample_without_replacement():
     """sample_without_replacement with env_ids should return one layout per env from correct variant."""
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
@@ -315,6 +377,8 @@ def test_pooled_placer_heterogeneous_sample_without_replacement():
 
 def test_pooled_placer_heterogeneous_sample_without_replacement_requires_env_ids():
     """Heterogeneous pool should assert when env_ids is not provided."""
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
@@ -324,6 +388,8 @@ def test_pooled_placer_heterogeneous_sample_without_replacement_requires_env_ids
 
 def test_pooled_placer_heterogeneous_sample_with_replacement():
     """sample_with_replacement should return per-variant layouts without consuming."""
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
@@ -335,6 +401,8 @@ def test_pooled_placer_heterogeneous_sample_with_replacement():
 
 def test_pooled_placer_heterogeneous_refill():
     """Exhausting a variant sub-pool should trigger a refill."""
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=4, num_envs=2)
 
@@ -352,6 +420,13 @@ def test_pooled_placer_heterogeneous_refill():
 
 def test_pooled_placer_homogeneous_unaffected_by_num_envs():
     """Homogeneous pool should work the same whether num_envs is passed or not."""
+    from isaaclab_arena.assets.dummy_object import DummyObject
+    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
+    from isaaclab_arena.relations.relations import On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
     desk = _make_desk()
     box = DummyObject(
         name="box",
@@ -379,6 +454,14 @@ def test_pooled_placer_multi_set_different_variant_counts():
     Bottles (3 variants) and boxes (2 variants) across 6 envs.
     Each env gets its own pool with layouts matching its object geometry.
     """
+    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
+    from isaaclab_arena.relations.relations import On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
+
     desk = _make_desk()
 
     bottle_small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.2))
@@ -409,6 +492,14 @@ def test_pooled_placer_multi_set_different_variant_counts():
 
 def test_pooled_placer_multi_set_sample_with_replacement():
     """sample_with_replacement with multi-set heterogeneous objects."""
+    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
+    from isaaclab_arena.relations.relations import On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
+
     desk = _make_desk()
 
     a_s = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.15))
@@ -435,6 +526,14 @@ def test_pooled_placer_multi_set_sample_with_replacement():
 
 def test_pooled_placer_multi_set_refill():
     """Exhausting a per-env pool should trigger refill with multi-set objects."""
+    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
+    from isaaclab_arena.relations.relations import On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
+
     desk = _make_desk()
 
     v1 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.15))
@@ -466,6 +565,8 @@ def test_pooled_placer_multi_set_refill():
 
 def test_pooled_placer_per_env_pools_isolated():
     """Each env_id should have its own independent pool of layouts."""
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
diff --git a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
index 280e9b909..f45b3f1ab 100644
--- a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
+++ b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
@@ -180,6 +180,7 @@ def _enable_periodic_reset(cfg):
             task=NoTask(),
             teleop_device=teleop_device,
             env_cfg_callback=env_cfg_callback,
+            force_convex_hull=(mode == "heterogeneous"),
         )
         return isaaclab_arena_environment
 
@@ -209,6 +210,10 @@ def _build_heterogeneous_objects(self, tabletop_reference, object_names=None):
         from isaaclab_arena.relations.relations import AtPosition, On
 
         if object_names:
+            print(
+                "Warning: --objects with --mode heterogeneous wraps each object as a "
+                "single-variant set (no per-env variance). Use default sets for true heterogeneity."
+            )
             placeable_assets = []
             for name in object_names:
                 obj = self.asset_registry.get_asset_by_name(name)()
@@ -241,7 +246,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> None:
             help=(
                 "Object names (works in both modes). "
                 f"Homo default: {' '.join(DEFAULT_TABLE_OBJECTS)}; "
-                "Hetero default: multi-variant sets (cracker_box/sugar_box, tomato_soup_can/mustard_bottle)"
+                f"Hetero default: {', '.join(HETERO_VARIANT_SETS.keys())} variant sets"
             ),
         )
         parser.add_argument("--embodiment", type=str, default="gr1_joint", help="Robot embodiment to use")

From a4c605d43a01b8d2a31ad19e85d1a837feb1f25d Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 18:27:30 -0700
Subject: [PATCH 15/46] test import

---
 .../tests/test_heterogeneous_placement.py     | 163 ++++--------------
 1 file changed, 31 insertions(+), 132 deletions(-)

diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 4ba0f9f79..8c39ae3be 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -9,48 +9,44 @@
 
 import pytest
 
+from isaaclab_arena.assets.dummy_object import DummyObject
+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.pooled_object_placer import PooledObjectPlacer
+from isaaclab_arena.relations.relation_solver import RelationSolver
+from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
+from isaaclab_arena.relations.relations import IsAnchor, On
+from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+from isaaclab_arena.utils.pose import Pose
+
 # ---------------------------------------------------------------------------
-# Helpers — all isaaclab_arena imports are deferred to avoid importing
-# Isaac Sim modules at pytest collection time.
+# Helpers
 # ---------------------------------------------------------------------------
 
 
-def _make_heterogeneous_dummy_class():
-    """Return the HeterogeneousDummyObject class (deferred import of DummyObject)."""
-
-    from isaaclab_arena.assets.dummy_object import DummyObject
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
-    class HeterogeneousDummyObject(DummyObject):
-        """DummyObject that provides different bounding boxes per environment.
-
-        Used to exercise the heterogeneous placement path without requiring
-        RigidObjectSet's USD machinery.
-        """
-
-        def __init__(self, name: str, bboxes: list[AxisAlignedBoundingBox], **kwargs):
-            super().__init__(name=name, bounding_box=bboxes[0], **kwargs)
-            self._per_env_bboxes = bboxes
-            self.heterogeneous_bbox = True
-            self.objects = bboxes
+class HeterogeneousDummyObject(DummyObject):
+    """DummyObject that provides different bounding boxes per environment.
 
-        def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
-            n_variants = len(self._per_env_bboxes)
-            indices = [i % n_variants for i in range(num_envs)]
-            min_pts = torch.stack([self._per_env_bboxes[idx].min_point[0] for idx in indices])
-            max_pts = torch.stack([self._per_env_bboxes[idx].max_point[0] for idx in indices])
-            return AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts)
+    Used to exercise the heterogeneous placement path without requiring
+    RigidObjectSet's USD machinery.
+    """
 
-    return HeterogeneousDummyObject
+    def __init__(self, name: str, bboxes: list[AxisAlignedBoundingBox], **kwargs):
+        super().__init__(name=name, bounding_box=bboxes[0], **kwargs)
+        self._per_env_bboxes = bboxes
+        self.heterogeneous_bbox = True
+        self.objects = bboxes
 
+    def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
+        n_variants = len(self._per_env_bboxes)
+        indices = [i % n_variants for i in range(num_envs)]
+        min_pts = torch.stack([self._per_env_bboxes[idx].min_point[0] for idx in indices])
+        max_pts = torch.stack([self._per_env_bboxes[idx].max_point[0] for idx in indices])
+        return AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts)
 
-def _make_desk():
-    """Create a desk anchor object for tests."""
-    from isaaclab_arena.assets.dummy_object import DummyObject
-    from isaaclab_arena.relations.relations import IsAnchor
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-    from isaaclab_arena.utils.pose import Pose
 
+def _make_desk() -> DummyObject:
     desk = DummyObject(
         name="desk",
         bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)),
@@ -68,9 +64,6 @@ def _make_desk():
 def test_dummy_object_bbox_per_env_expands_single():
     """Default get_bounding_box_per_env should repeat the single bbox."""
 
-    from isaaclab_arena.assets.dummy_object import DummyObject
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
     obj = DummyObject(
         name="box",
         bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)),
@@ -85,10 +78,6 @@ def test_dummy_object_bbox_per_env_expands_single():
 def test_heterogeneous_dummy_returns_different_bboxes():
     """HeterogeneousDummyObject should cycle through its member bboxes."""
 
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
-    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
-
     small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
     large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3))
     obj = HeterogeneousDummyObject(name="set", bboxes=[small, large])
@@ -108,12 +97,6 @@ def test_heterogeneous_dummy_returns_different_bboxes():
 def test_solver_accepts_env_bboxes():
     """Solver should accept env_bboxes and produce valid results."""
 
-    from isaaclab_arena.assets.dummy_object import DummyObject
-    from isaaclab_arena.relations.relation_solver import RelationSolver
-    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
-    from isaaclab_arena.relations.relations import On
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
     desk = _make_desk()
     box = DummyObject(
         name="box",
@@ -149,15 +132,6 @@ def test_solver_accepts_env_bboxes():
 def test_placer_heterogeneous_produces_per_env_results():
     """Placer should detect heterogeneous objects and solve per-env."""
 
-    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.relation_solver_params import RelationSolverParams
-    from isaaclab_arena.relations.relations import On
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
-    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
-
     desk = _make_desk()
 
     small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
@@ -187,15 +161,6 @@ def test_placer_heterogeneous_produces_per_env_results():
 def test_placer_heterogeneous_z_height_matches_variant():
     """Objects should be placed at z-height matching their env's variant bbox."""
 
-    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.relation_solver_params import RelationSolverParams
-    from isaaclab_arena.relations.relations import On
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
-    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
-
     desk = _make_desk()
 
     # "tall" variant: height 0.4 → bottom at z ≈ 0.11 (desk top 0.1 + clearance 0.01)
@@ -234,16 +199,6 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
     bboxes per env while X is identical everywhere.
     """
 
-    from isaaclab_arena.assets.dummy_object import DummyObject
-    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.relation_solver_params import RelationSolverParams
-    from isaaclab_arena.relations.relations import On
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
-    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
-
     desk = _make_desk()
 
     # A: heterogeneous — small variant in even envs, large in odd envs.
@@ -259,6 +214,8 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
     )
     obj_x.add_relation(On(desk, clearance_m=0.01))
 
+    # No-overlap is handled automatically by the solver's built-in clearance.
+
     objects = [desk, obj_a, obj_x]
     num_envs = 4
 
@@ -296,14 +253,6 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
 def test_homogeneous_path_unchanged():
     """When no heterogeneous objects exist, the homogeneous path is used."""
 
-    from isaaclab_arena.assets.dummy_object import DummyObject
-    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.relation_solver_params import RelationSolverParams
-    from isaaclab_arena.relations.relations import On
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
     desk = _make_desk()
     box = DummyObject(
         name="box",
@@ -332,13 +281,6 @@ def test_homogeneous_path_unchanged():
 
 def _make_hetero_pool_objects():
     """Create desk + heterogeneous box for pool tests."""
-    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
-    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
-    from isaaclab_arena.relations.relations import On
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
-    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
-
     desk = _make_desk()
     small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
     large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3))
@@ -352,8 +294,6 @@ def _make_hetero_pool_objects():
 
 def test_pooled_placer_heterogeneous_is_detected():
     """PooledObjectPlacer should detect heterogeneous objects and create variant sub-pools."""
-    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
-
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
@@ -363,8 +303,6 @@ def test_pooled_placer_heterogeneous_is_detected():
 
 def test_pooled_placer_heterogeneous_sample_without_replacement():
     """sample_without_replacement with env_ids should return one layout per env from correct variant."""
-    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
-
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
@@ -377,8 +315,6 @@ def test_pooled_placer_heterogeneous_sample_without_replacement():
 
 def test_pooled_placer_heterogeneous_sample_without_replacement_requires_env_ids():
     """Heterogeneous pool should assert when env_ids is not provided."""
-    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
-
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
@@ -388,8 +324,6 @@ def test_pooled_placer_heterogeneous_sample_without_replacement_requires_env_ids
 
 def test_pooled_placer_heterogeneous_sample_with_replacement():
     """sample_with_replacement should return per-variant layouts without consuming."""
-    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
-
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
@@ -401,8 +335,6 @@ def test_pooled_placer_heterogeneous_sample_with_replacement():
 
 def test_pooled_placer_heterogeneous_refill():
     """Exhausting a variant sub-pool should trigger a refill."""
-    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
-
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=4, num_envs=2)
 
@@ -420,13 +352,6 @@ def test_pooled_placer_heterogeneous_refill():
 
 def test_pooled_placer_homogeneous_unaffected_by_num_envs():
     """Homogeneous pool should work the same whether num_envs is passed or not."""
-    from isaaclab_arena.assets.dummy_object import DummyObject
-    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
-    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
-    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
-    from isaaclab_arena.relations.relations import On
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
     desk = _make_desk()
     box = DummyObject(
         name="box",
@@ -454,14 +379,6 @@ def test_pooled_placer_multi_set_different_variant_counts():
     Bottles (3 variants) and boxes (2 variants) across 6 envs.
     Each env gets its own pool with layouts matching its object geometry.
     """
-    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
-    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
-    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
-    from isaaclab_arena.relations.relations import On
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
-    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
-
     desk = _make_desk()
 
     bottle_small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.2))
@@ -492,14 +409,6 @@ def test_pooled_placer_multi_set_different_variant_counts():
 
 def test_pooled_placer_multi_set_sample_with_replacement():
     """sample_with_replacement with multi-set heterogeneous objects."""
-    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
-    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
-    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
-    from isaaclab_arena.relations.relations import On
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
-    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
-
     desk = _make_desk()
 
     a_s = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.08, 0.08, 0.15))
@@ -526,14 +435,6 @@ def test_pooled_placer_multi_set_sample_with_replacement():
 
 def test_pooled_placer_multi_set_refill():
     """Exhausting a per-env pool should trigger refill with multi-set objects."""
-    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
-    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
-    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
-    from isaaclab_arena.relations.relations import On
-    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
-    HeterogeneousDummyObject = _make_heterogeneous_dummy_class()
-
     desk = _make_desk()
 
     v1 = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.15))
@@ -565,8 +466,6 @@ def test_pooled_placer_multi_set_refill():
 
 def test_pooled_placer_per_env_pools_isolated():
     """Each env_id should have its own independent pool of layouts."""
-    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
-
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 

From 62e7673c2d0d85cc656491f85bac27fbf78cc695 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 12:10:31 -0700
Subject: [PATCH 16/46] remove dependency of convex hull

---
 isaaclab_arena/assets/background_library.py   | 20 ++++++++++++++++
 .../environments/arena_env_builder.py         | 23 -------------------
 .../isaaclab_arena_environment.py             |  8 +------
 ...e_multi_object_no_collision_environment.py | 21 ++++++++---------
 4 files changed, 31 insertions(+), 41 deletions(-)

diff --git a/isaaclab_arena/assets/background_library.py b/isaaclab_arena/assets/background_library.py
index decceadd6..df61285a3 100644
--- a/isaaclab_arena/assets/background_library.py
+++ b/isaaclab_arena/assets/background_library.py
@@ -5,6 +5,7 @@
 
 from typing import Any
 
+import isaaclab.sim as sim_utils
 from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR
 
 from isaaclab_arena.assets.background import Background
@@ -139,6 +140,25 @@ def __init__(self):
         super().__init__()
 
 
+@register_asset
+class OfficeTableBackground(LibraryBackground):
+    """
+    A basic office table.
+    """
+
+    name = "office_table_background"
+    tags = ["background"]
+    usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Mimic/nut_pour_task/nut_pour_assets/table.usd"
+    object_min_z = -0.05
+    scale = (1.0, 1.0, 0.7)
+    spawn_cfg_addon = {
+        "rigid_props": sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True),
+    }
+
+    def __init__(self):
+        super().__init__(scale=self.scale)
+
+
 @register_asset
 class LightwheelKitchenBackground(LibraryBackground):
     """
diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py
index bb74cc8cd..f0f4a5b8f 100644
--- a/isaaclab_arena/environments/arena_env_builder.py
+++ b/isaaclab_arena/environments/arena_env_builder.py
@@ -292,30 +292,7 @@ def make_registered_and_return_cfg(
     ) -> tuple[ManagerBasedEnv, IsaacLabArenaManagerBasedRLEnvCfg]:
         name, cfg = self.build_registered(env_cfg)
         env = gym.make(name, cfg=cfg, render_mode=render_mode)
-        if self.arena_env.force_convex_hull:
-            _force_convex_hull(env)
         # ViewportCameraController sets the camera before KitVisualizer.initialize() is called,
         # so the call is silently ignored. Re-apply here once the visualizers are fully initialized.
         reapply_viewer_cfg(env)
         return env, cfg
-
-
-def _force_convex_hull(env: ManagerBasedEnv) -> None:
-    """Replace ``convexDecomposition`` with ``convexHull`` on all MeshCollision prims.
-
-    ``convexDecomposition`` on raw scanned meshes (e.g. robolab assets) creates
-    irregular contact surfaces that are unstable in multi-object scenarios.
-    ``convexHull`` produces a single convex shape that behaves predictably.
-    """
-    from pxr import UsdPhysics
-
-    stage = env.unwrapped.sim.stage
-    for prim in stage.Traverse():
-        if not prim.HasAPI(UsdPhysics.MeshCollisionAPI):
-            continue
-        mesh_col = UsdPhysics.MeshCollisionAPI(prim)
-        approx_attr = mesh_col.GetApproximationAttr()
-        if not approx_attr or not approx_attr.HasValue():
-            continue
-        if approx_attr.Get() == "convexDecomposition":
-            approx_attr.Set("convexHull")
diff --git a/isaaclab_arena/environments/isaaclab_arena_environment.py b/isaaclab_arena/environments/isaaclab_arena_environment.py
index 8498d5a56..9ed801ff2 100644
--- a/isaaclab_arena/environments/isaaclab_arena_environment.py
+++ b/isaaclab_arena/environments/isaaclab_arena_environment.py
@@ -26,10 +26,9 @@ def __init__(
         embodiment: EmbodimentBase | None = None,
         task: TaskBase | None = None,
         teleop_device: TeleopDeviceBase | None = None,
-        env_cfg_callback: Callable[IsaacLabArenaManagerBasedRLEnvCfg] | None = None,
+        env_cfg_callback: Callable[[IsaacLabArenaManagerBasedRLEnvCfg], IsaacLabArenaManagerBasedRLEnvCfg] | None = None,
         rl_framework_entry_point: str | None = None,
         rl_policy_cfg: str | None = None,
-        force_convex_hull: bool = False,
     ):
         """
         Args:
@@ -47,10 +46,6 @@ def __init__(
                 ``rl_policy_cfg`` is set.
             rl_policy_cfg: Import path to the RL policy config class, e.g.
                 ``"my_module:RLPolicyCfg"``.
-            force_convex_hull: If True, replace ``convexDecomposition`` with ``convexHull``
-                on all MeshCollision prims after scene creation. Needed for assets with
-                raw scanned meshes (e.g. robolab objects) that are unstable with
-                ``convexDecomposition``.
         """
         self.name = name
         self.scene = scene
@@ -58,7 +53,6 @@ def __init__(
         self.task = task
         self.teleop_device = teleop_device
         self.env_cfg_callback = env_cfg_callback
-        self.force_convex_hull = force_convex_hull
         if (rl_framework_entry_point is None) != (rl_policy_cfg is None):
             raise ValueError("rl_framework_entry_point and rl_policy_cfg must both be set or both be None.")
         self.rl_framework_entry_point = rl_framework_entry_point
diff --git a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
index f45b3f1ab..c324cc7e3 100644
--- a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
+++ b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
@@ -76,10 +76,10 @@
         "measuring_spoon_handal_robolab",
     ],
     "boxes": [
-        "popcorn_box_hope_robolab",
-        "chocolate_pudding_mix_hope_robolab",
-        "macaroni_and_cheese_hope_robolab",
-        "granola_bars_hope_robolab",
+        "butter_hope_robolab",
+        "raisin_box_hope_robolab",
+        "yogurt_cup_hope_robolab",
+        "oatmeal_raisin_cookies_hope_robolab",
     ],
 }
 
@@ -130,10 +130,10 @@ def get_env(self, args_cli: argparse.Namespace) -> IsaacLabArenaEnvironment:
         ground_plane = self.asset_registry.get_asset_by_name("ground_plane")()
         light = self.asset_registry.get_asset_by_name("light")()
 
-        table_background = self.asset_registry.get_asset_by_name("office_table")()
+        table_background = self.asset_registry.get_asset_by_name("office_table_background")()
         tabletop_reference = ObjectReference(
             name="table",
-            prim_path="{ENV_REGEX_NS}/office_table/Geometry/sm_tabletop_a01_01/sm_tabletop_a01_top_01",
+            prim_path="{ENV_REGEX_NS}/office_table_background/Geometry/sm_tabletop_a01_01/sm_tabletop_a01_top_01",
             parent_asset=table_background,
         )
         tabletop_reference.add_relation(IsAnchor())
@@ -180,7 +180,6 @@ def _enable_periodic_reset(cfg):
             task=NoTask(),
             teleop_device=teleop_device,
             env_cfg_callback=env_cfg_callback,
-            force_convex_hull=(mode == "heterogeneous"),
         )
         return isaaclab_arena_environment
 
@@ -195,7 +194,7 @@ def _build_homogeneous_objects(self, tabletop_reference, object_names=None):
         placeable_assets = []
         for name in names:
             obj = self.asset_registry.get_asset_by_name(name)()
-            obj.add_relation(On(tabletop_reference, clearance_m=0.001))
+            obj.add_relation(On(tabletop_reference, clearance_m=0.01))
             placeable_assets.append(obj)
         return placeable_assets
 
@@ -218,20 +217,20 @@ def _build_heterogeneous_objects(self, tabletop_reference, object_names=None):
             for name in object_names:
                 obj = self.asset_registry.get_asset_by_name(name)()
                 obj_set = RigidObjectSet(name=name, objects=[obj])
-                obj_set.add_relation(On(tabletop_reference, clearance_m=0.001))
+                obj_set.add_relation(On(tabletop_reference, clearance_m=0.01))
                 placeable_assets.append(obj_set)
         else:
             placeable_assets = []
             for name, x, y in HETERO_FIXED_OBJECTS:
                 obj = self.asset_registry.get_asset_by_name(name)()
-                obj.add_relation(On(tabletop_reference, clearance_m=0.001))
+                obj.add_relation(On(tabletop_reference, clearance_m=0.01))
                 obj.add_relation(AtPosition(x=x, y=y))
                 placeable_assets.append(obj)
 
             for set_name, variant_names in HETERO_VARIANT_SETS.items():
                 members = [self.asset_registry.get_asset_by_name(n)() for n in variant_names]
                 obj_set = RigidObjectSet(name=set_name, objects=members)
-                obj_set.add_relation(On(tabletop_reference, clearance_m=0.001))
+                obj_set.add_relation(On(tabletop_reference, clearance_m=0.01))
                 placeable_assets.append(obj_set)
 
         return placeable_assets

From a72a41aee6d970a94e29612152fd4af5199b4d6e Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 14:40:50 -0700
Subject: [PATCH 17/46] clean pool strategies

---
 .../isaaclab_arena_environment.py             |   4 +-
 isaaclab_arena/relations/placement_events.py  |  12 +-
 .../relations/pooled_object_placer.py         | 289 ++++++++----------
 .../tests/test_heterogeneous_placement.py     | 125 +++++---
 isaaclab_arena/tests/test_placement_events.py |   4 +-
 5 files changed, 225 insertions(+), 209 deletions(-)

diff --git a/isaaclab_arena/environments/isaaclab_arena_environment.py b/isaaclab_arena/environments/isaaclab_arena_environment.py
index 9ed801ff2..b481cab0a 100644
--- a/isaaclab_arena/environments/isaaclab_arena_environment.py
+++ b/isaaclab_arena/environments/isaaclab_arena_environment.py
@@ -26,7 +26,9 @@ def __init__(
         embodiment: EmbodimentBase | None = None,
         task: TaskBase | None = None,
         teleop_device: TeleopDeviceBase | None = None,
-        env_cfg_callback: Callable[[IsaacLabArenaManagerBasedRLEnvCfg], IsaacLabArenaManagerBasedRLEnvCfg] | None = None,
+        env_cfg_callback: (
+            Callable[[IsaacLabArenaManagerBasedRLEnvCfg], IsaacLabArenaManagerBasedRLEnvCfg] | None
+        ) = None,
         rl_framework_entry_point: str | None = None,
         rl_policy_cfg: str | None = None,
     ):
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index ff0545b41..7eb47044c 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -34,8 +34,9 @@ def solve_and_place_objects(
 ) -> None:
     """Coordinated reset event that draws layouts from the pool and writes poses.
 
-    Registered as a single ``EventTermCfg(mode="reset")``. Each call draws one
-    layout per resetting environment from the pool and writes the poses to sim.
+    Registered as a single ``EventTermCfg(mode="reset")``. Each call advances
+    the placement pool by one full env round, then writes poses only for the
+    environments being reset.
 
     Args:
         env: The Isaac Lab environment.
@@ -46,16 +47,15 @@ def solve_and_place_objects(
     if env_ids is None or len(env_ids) == 0:
         return
 
-    num_reset_envs = len(env_ids)
-    layouts_per_env = placement_pool.sample_without_replacement(num_reset_envs, env_ids=env_ids)
+    all_results = placement_pool.sample_without_replacement(env.scene.env_origins.shape[0])
 
     anchor_objects_set = set(get_anchor_objects(objects))
     rotations = {obj: get_rotation_xyzw(obj) for obj in objects if obj not in anchor_objects_set}
 
     zero_velocity = torch.zeros(1, 6, device=env.device)
-    for local_idx, cur_env in enumerate(env_ids.tolist()):
+    for cur_env in env_ids.tolist():
         env_id_tensor = torch.tensor([cur_env], device=env.device)
-        positions = layouts_per_env[local_idx].positions
+        positions = all_results[cur_env].positions
         for obj, pos in positions.items():
             if obj in anchor_objects_set:
                 continue
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index af36d879e..21c434362 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -24,19 +24,12 @@ class PooledObjectPlacer:
     ``pool_size``, keeping only those that pass validation. The pool is refilled
     automatically when consumed layouts run out.
 
-    **Homogeneous mode** (default): all objects have the same geometry in
-    every environment.  Layouts are stored in a single flat list and any
-    layout can serve any environment.
-
-    **Heterogeneous mode** (activated when any object has
-    ``heterogeneous_bbox = True``, e.g. ``RigidObjectSet``): each
-    environment has its own fixed set of object variants, assigned at
-    build time.  Layouts are stored per ``env_id`` so that resets always
-    return a layout solved for that environment's actual object geometry.
+    Layouts are always stored in per-env pools.  The public sampling methods
+    expose only the replacement strategy; internally, samples are drawn in
+    env-index order.
 
     * :meth:`sample_without_replacement` — returns the next *count* layouts
-      sequentially.  Auto-refills when exhausted.  In heterogeneous mode,
-      pass ``env_ids`` so each environment receives a matching layout.
+      sequentially.  Auto-refills when exhausted.
     * :meth:`sample_with_replacement` — picks *count* layouts at random
       (non-consuming).  Used for static initial positions.
 
@@ -44,8 +37,8 @@ class PooledObjectPlacer:
         objects: All objects (including anchors) participating in relation solving.
         placer_params: Parameters forwarded to ``ObjectPlacer`` for the batched solve.
         pool_size: Number of layouts to solve per batch.
-        num_envs: Total number of simulation environments.  Required for
-            heterogeneous placement so per-env pools can be created.
+        num_envs: Total number of simulation environments.  Required when
+            layouts use env-specific object variants and defaults to 1 otherwise.
     """
 
     def __init__(
@@ -58,62 +51,82 @@ def __init__(
         if pool_size < 1:
             raise ValueError(f"pool_size must be >= 1, got {pool_size}")
 
-        self._objects = objects
+        self._objects = list(objects)
         self._placer = ObjectPlacer(params=placer_params)
         self._pool_size = pool_size
-        self._heterogeneous = any(getattr(obj, "heterogeneous_bbox", False) for obj in objects)
-
-        if self._heterogeneous:
-            assert (
-                num_envs is not None
-            ), "num_envs is required for heterogeneous placement so per-env pools can be created."
-            self._num_envs = num_envs
-
-            self._layout_pools: dict[int, list[PlacementResult]] = {env_id: [] for env_id in range(num_envs)}
-            self._layout_cursors: dict[int, int] = {env_id: 0 for env_id in range(num_envs)}
-
-            self._solve_and_store_heterogeneous(pool_size)
-            for env_id, pool in self._layout_pools.items():
-                if not pool:
-                    raise RuntimeError(
-                        f"Placement pool failed to produce any valid layouts for env {env_id} "
-                        f"from {pool_size} attempts. Check object relations and constraints."
-                    )
-        else:
-            self._layouts: list[PlacementResult] = []
-            self._next_idx: int = 0
-
-            # Pre-solve the initial batch (runs the gradient solver, no simulation is needed).
-            self._solve_and_store(pool_size)
-            if not self._layouts:
+        self._uses_env_specific_bboxes = any(getattr(obj, "heterogeneous_bbox", False) for obj in objects)
+
+        self._num_envs = num_envs if num_envs is not None else 1
+        if self._num_envs < 1:
+            raise ValueError(f"num_envs must be >= 1, got {self._num_envs}")
+        if self._uses_env_specific_bboxes:
+            assert num_envs is not None, "num_envs is required when layouts use env-specific object variants."
+        self._layout_pools: dict[int, list[PlacementResult]] = {cur_env: [] for cur_env in range(self._num_envs)}
+        self._layout_cursors: dict[int, int] = {cur_env: 0 for cur_env in range(self._num_envs)}
+
+        self._solve_and_store(pool_size)
+        for cur_env, pool in self._layout_pools.items():
+            if not pool:
                 raise RuntimeError(
-                    f"Pooled object placer failed to produce any valid layouts from {pool_size} attempts. "
-                    "Check object relations and constraints."
+                    f"Placement pool failed to produce any valid layouts for env {cur_env} "
+                    f"from {pool_size} attempts. Check object relations and constraints."
                 )
 
-    @property
-    def is_heterogeneous(self) -> bool:
-        """Whether this pool operates in heterogeneous (per-env) mode."""
-        return self._heterogeneous
-
     # ------------------------------------------------------------------
-    # Homogeneous (flat pool) internals
+    # Pool storage internals
     # ------------------------------------------------------------------
 
-    def _compact(self) -> None:
-        """Drop consumed layouts and reset the read index to free memory."""
-        self._layouts = self._layouts[self._next_idx :]
-        self._next_idx = 0
+    def _discard_consumed_layouts(self) -> None:
+        """Drop consumed layouts from every env pool before appending new layouts."""
+        for cur_env in self._layout_pools:
+            idx = self._layout_cursors[cur_env]
+            self._layout_pools[cur_env] = self._layout_pools[cur_env][idx:]
+            self._layout_cursors[cur_env] = 0
 
     def _solve_and_store(self, num_layouts: int) -> None:
-        """Solve *num_layouts* placements and append valid layouts to the pool.
+        """Solve layouts and append complete per-env rounds to the pools."""
+        self._discard_consumed_layouts()
+        target_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
+        max_solve_batches = max(1, self._placer.params.max_placement_attempts)
+
+        for _ in range(max_solve_batches):
+            missing_per_env = [
+                target_per_env - (len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env])
+                for cur_env in range(self._num_envs)
+            ]
+            max_missing = max(missing_per_env)
+            if max_missing <= 0:
+                return
+
+            batch_size = max_missing * self._num_envs
+            if self._uses_env_specific_bboxes:
+                all_results, layouts_per_env = self._solve_layouts_with_env_bboxes(batch_size)
+                self._store_env_matched_results(all_results, layouts_per_env)
+            else:
+                layouts = self._solve_reusable_layouts(batch_size)
+                self._store_reusable_results(layouts)
+
+            available = [
+                len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] for cur_env in range(self._num_envs)
+            ]
+            if min(available) >= target_per_env:
+                return
+
+        available = [
+            len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] for cur_env in range(self._num_envs)
+        ]
+        raise RuntimeError(
+            f"Placement pool could not fill {target_per_env} layouts per env after "
+            f"{max_solve_batches} solve batches. Available per env: {available}."
+        )
+
+    def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
+        """Solve layouts that can be used by any env pool.
 
         When no layouts pass strict validation, the best-loss layouts are
         accepted with a warning (matching pre-pool behaviour where validation
         failures were non-fatal).
         """
-        self._compact()
-
         with torch.inference_mode(False):
             result = self._placer.place(self._objects, num_envs=num_layouts, result_per_env=True)
 
@@ -128,34 +141,25 @@ def _solve_and_store(self, num_layouts: int) -> None:
             )
 
         if valid_layouts:
-            self._layouts.extend(valid_layouts)
-        else:
-            print("Warning: No layouts passed strict validation. Accepting best-loss layouts as fallback.")
-            self._layouts.extend(all_layouts)
-
-    def _refill_pool_via_solve_if_required(self, count: int) -> None:
-        """Refill the pool via solve when fewer than *count* layouts are available."""
-        if self.remaining < count:
-            self._solve_and_store(max(self._pool_size, count))
+            return valid_layouts
 
-        if self.remaining < count:
-            raise RuntimeError(
-                f"Pooled object placer has {self.remaining} valid layouts but {count} were requested. "
-                "The solver is not producing enough valid placements."
-            )
+        print("Warning: No candidates passed strict validation. Accepting best-loss layouts as fallback.")
+        return all_layouts
 
-    # ------------------------------------------------------------------
-    # Heterogeneous (per-env pool) internals
-    # ------------------------------------------------------------------
+    def _store_reusable_results(self, layouts: list[PlacementResult]) -> None:
+        """Distribute reusable layouts across env pools without dropping valid results."""
+        if not layouts:
+            return
 
-    def _compact_env_pool(self, env_id: int) -> None:
-        """Drop consumed layouts for a single env and reset its cursor."""
-        idx = self._layout_cursors[env_id]
-        self._layout_pools[env_id] = self._layout_pools[env_id][idx:]
-        self._layout_cursors[env_id] = 0
+        for layout in layouts:
+            cur_env = min(
+                range(self._num_envs),
+                key=lambda cur_env: len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env],
+            )
+            self._layout_pools[cur_env].append(layout)
 
-    def _solve_and_store_heterogeneous(self, num_layouts: int) -> None:
-        """Solve layouts and store valid results into per-env pools.
+    def _solve_layouts_with_env_bboxes(self, num_layouts: int) -> tuple[list[PlacementResult], int]:
+        """Solve layouts tied to each env's actual object geometry.
 
         Computes bounding boxes for the real ``num_envs`` once, tiles them
         to ``num_layouts`` entries, and solves everything in **one** batched
@@ -164,17 +168,14 @@ def _solve_and_store_heterogeneous(self, num_layouts: int) -> None:
         """
         from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
-        for env_id in self._layout_pools:
-            self._compact_env_pool(env_id)
-
-        layouts_per_env = max(1, num_layouts // self._num_envs)
+        layouts_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
         total_layouts = layouts_per_env * self._num_envs
 
         real_bboxes: dict = {obj: obj.get_bounding_box_per_env(self._num_envs) for obj in self._objects}
 
         tiled_bboxes: dict = {}
         for obj, bbox in real_bboxes.items():
-            # (num_envs, 3) → repeat each env's row layouts_per_env times → (total_layouts, 3)
+            # (num_envs, 3) -> repeat each env's row layouts_per_env times -> (total_layouts, 3)
             min_pt = bbox.min_point.repeat_interleave(layouts_per_env, dim=0)
             max_pt = bbox.max_point.repeat_interleave(layouts_per_env, dim=0)
             tiled_bboxes[obj] = AxisAlignedBoundingBox(min_point=min_pt, max_point=max_pt)
@@ -188,134 +189,103 @@ def _solve_and_store_heterogeneous(self, num_layouts: int) -> None:
             )
 
         all_results = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
+        return all_results, layouts_per_env
 
+    def _store_env_matched_results(self, all_results: list[PlacementResult], layouts_per_env: int) -> None:
+        """Store layouts into the env pools they were solved for."""
         total_valid = 0
         for i, r in enumerate(all_results):
-            env_id = i // layouts_per_env
+            cur_env = i // layouts_per_env
             if r.success:
-                self._layout_pools[env_id].append(r)
+                self._layout_pools[cur_env].append(r)
                 total_valid += 1
 
         total_solved = len(all_results)
         if total_valid < total_solved:
-            failed_envs = [e for e in self._layout_pools if not self._layout_pools[e]]
+            failed_envs = [cur_env for cur_env in self._layout_pools if not self._layout_pools[cur_env]]
             msg = (
-                f"Placement pool (heterogeneous): solved {total_solved} candidates,"
+                f"Placement pool (env-matched layouts): solved {total_solved} candidates,"
                 f" {total_valid} valid, {total_solved - total_valid} failed validation"
             )
             if failed_envs:
                 msg += f". Envs with zero valid layouts: {failed_envs}"
             print(msg)
 
-        for env_id in range(self._num_envs):
-            if self._layout_pools[env_id]:
-                continue
+        for cur_env in range(self._num_envs):
             best: PlacementResult | None = None
-            start = env_id * layouts_per_env
+            had_valid = False
+            start = cur_env * layouts_per_env
             end = start + layouts_per_env
             for r in all_results[start:end]:
+                if r.success:
+                    had_valid = True
+                    continue
                 if best is None or r.final_loss < best.final_loss:
                     best = r
-            if best is not None:
+            if not had_valid and best is not None:
                 print(
-                    f"Warning: env {env_id} had no valid layouts; "
+                    f"Warning: env {cur_env} had too few valid layouts; "
                     f"accepting best-loss fallback (loss={best.final_loss:.6f})."
                 )
-                self._layout_pools[env_id].append(best)
+                self._layout_pools[cur_env].append(best)
 
     # ------------------------------------------------------------------
     # Public API
     # ------------------------------------------------------------------
 
-    def sample_without_replacement(
-        self, count: int, env_ids: list[int] | torch.Tensor | None = None
-    ) -> list[PlacementResult]:
+    def sample_without_replacement(self, count: int) -> list[PlacementResult]:
         """Return the next *count* layouts sequentially (without replacement).
 
-        Auto-refills the pool when there are not enough layouts ahead of the
-        read index.
-
-        In **heterogeneous mode** ``env_ids`` must be provided so each
-        environment receives a layout matching its object geometry.
+        Auto-refills any env pool that does not have enough layouts ahead of
+        its read cursor.
 
         Args:
             count: Number of layouts to return.
-            env_ids: Environment indices being reset.  Required when the
-                pool is heterogeneous; ignored otherwise.
 
         Raises:
+            ValueError: If *count* is not a complete env round.
             RuntimeError: If the pool cannot provide *count* layouts after refilling.
         """
-        if self._heterogeneous:
-            return self._sample_without_replacement_heterogeneous(count, env_ids)
-
-        self._refill_pool_via_solve_if_required(count)
-
-        start = self._next_idx
-        self._next_idx += count
-        layouts = self._layouts[start : self._next_idx]
-        return layouts
-
-    def _sample_without_replacement_heterogeneous(
-        self, count: int, env_ids: list[int] | torch.Tensor | None
-    ) -> list[PlacementResult]:
-        """Draw one layout per requested env, refilling any depleted per-env pool."""
-        assert env_ids is not None, "env_ids must be provided for heterogeneous placement pools."
-
-        if isinstance(env_ids, torch.Tensor):
-            ids: list[int] = [int(x) for x in env_ids]
-        else:
-            ids = list(env_ids)
-        assert len(ids) == count
-
-        # Refill any env pool that doesn't have enough layouts.
-        demand_per_env: dict[int, int] = {}
-        for env_id in ids:
-            demand_per_env[env_id] = demand_per_env.get(env_id, 0) + 1
-
-        needs_refill = False
-        for env_id, demand in demand_per_env.items():
-            available = len(self._layout_pools[env_id]) - self._layout_cursors[env_id]
-            if available < demand:
-                needs_refill = True
-                break
+        if count % self._num_envs != 0:
+            raise ValueError(f"count must be a multiple of num_envs ({self._num_envs}), got {count}")
+
+        sample_env_order = [i % self._num_envs for i in range(count)]
+        layouts_per_env = count // self._num_envs
+        needs_refill = any(
+            len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] < layouts_per_env
+            for cur_env in range(self._num_envs)
+        )
 
         if needs_refill:
-            max_demand = max(demand_per_env.values())
-            self._solve_and_store_heterogeneous(max(self._pool_size, max_demand * self._num_envs))
+            self._solve_and_store(max(self._pool_size, count))
 
         results: list[PlacementResult] = []
-        for env_id in ids:
-            idx = self._layout_cursors[env_id]
-            if idx >= len(self._layout_pools[env_id]):
+        for cur_env in sample_env_order:
+            idx = self._layout_cursors[cur_env]
+            if idx >= len(self._layout_pools[cur_env]):
                 raise RuntimeError(
-                    f"Placement pool: env {env_id} has no more valid layouts. "
+                    f"Placement pool: env {cur_env} has no more valid layouts. "
                     "The solver is not producing enough valid placements."
                 )
-            results.append(self._layout_pools[env_id][idx])
-            self._layout_cursors[env_id] = idx + 1
+            results.append(self._layout_pools[cur_env][idx])
+            self._layout_cursors[cur_env] = idx + 1
 
         return results
 
     def sample_with_replacement(self, count: int) -> list[PlacementResult]:
         """Pick *count* layouts at random with replacement (non-consuming).
 
-        In **heterogeneous mode**, each position ``i`` in the returned
-        list corresponds to env ``i`` and is drawn from that env's pool.
+        Each returned layout is drawn from the per-env pool corresponding to
+        its position in the requested batch.
 
         Used by ``resolve_on_reset=False`` to assign initial positions
         that persist across resets.
         """
-        if self._heterogeneous:
-            return self._sample_with_replacement_heterogeneous(count)
-        return random.choices(self._layouts, k=count)
-
-    def _sample_with_replacement_heterogeneous(self, count: int) -> list[PlacementResult]:
-        """Pick one random layout per env from its pool (non-consuming)."""
+        sample_env_order = [i % self._num_envs for i in range(count)]
         results: list[PlacementResult] = []
-        for env_id in range(count):
-            pool = self._layout_pools[env_id]
-            assert pool, f"Env {env_id} has no valid layouts to sample from."
+        for cur_env in sample_env_order:
+            pool = self._layout_pools[cur_env]
+            assert pool, f"Env {cur_env} has no valid layouts to sample from."
             results.append(random.choice(pool))
         return results
 
@@ -323,11 +293,10 @@ def _sample_with_replacement_heterogeneous(self, count: int) -> list[PlacementRe
     def remaining(self) -> int:
         """Number of layouts not yet consumed by :meth:`sample_without_replacement`.
 
-        For heterogeneous pools, returns the minimum across all envs.
+        Reports the minimum available count across env pools so every env has
+        the same without-replacement capacity.
         """
-        if self._heterogeneous:
-            return min(len(self._layout_pools[e]) - self._layout_cursors[e] for e in self._layout_pools)
-        return len(self._layouts) - self._next_idx
+        return min(len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] for cur_env in self._layout_pools)
 
     @property
     def pool_size(self) -> int:
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 8c39ae3be..249024d8e 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -12,7 +12,7 @@
 from isaaclab_arena.assets.dummy_object import DummyObject
 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.placement_result import MultiEnvPlacementResult, PlacementResult
 from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
 from isaaclab_arena.relations.relation_solver import RelationSolver
 from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
@@ -275,7 +275,7 @@ def test_homogeneous_path_unchanged():
 
 
 # ---------------------------------------------------------------------------
-# PooledObjectPlacer heterogeneous mode
+# PooledObjectPlacer env-specific variants
 # ---------------------------------------------------------------------------
 
 
@@ -292,34 +292,36 @@ def _make_hetero_pool_objects():
     return desk, hetero, placer_params
 
 
-def test_pooled_placer_heterogeneous_is_detected():
-    """PooledObjectPlacer should detect heterogeneous objects and create variant sub-pools."""
+def test_pooled_placer_env_specific_layouts_sample_from_fixed_env_order():
+    """PooledObjectPlacer should hide env routing behind the sampling strategy."""
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
-    assert pool.is_heterogeneous
     assert pool.remaining > 0
+    draws = pool.sample_without_replacement(4)
+    assert len(draws) == 4
+    for draw in draws:
+        assert hetero in draw.positions
 
 
 def test_pooled_placer_heterogeneous_sample_without_replacement():
-    """sample_without_replacement with env_ids should return one layout per env from correct variant."""
+    """sample_without_replacement should return one layout per requested sample."""
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
-    env_ids = torch.tensor([0, 1, 2, 3])
-    draws = pool.sample_without_replacement(4, env_ids=env_ids)
+    draws = pool.sample_without_replacement(4)
     assert len(draws) == 4
     for d in draws:
         assert hetero in d.positions
 
 
-def test_pooled_placer_heterogeneous_sample_without_replacement_requires_env_ids():
-    """Heterogeneous pool should assert when env_ids is not provided."""
+def test_pooled_placer_heterogeneous_sample_without_replacement_requires_complete_rounds():
+    """sample_without_replacement should consume complete env rounds."""
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
-    with pytest.raises(AssertionError):
-        pool.sample_without_replacement(2, env_ids=None)
+    with pytest.raises(ValueError):
+        pool.sample_without_replacement(2)
 
 
 def test_pooled_placer_heterogeneous_sample_with_replacement():
@@ -340,18 +342,16 @@ def test_pooled_placer_heterogeneous_refill():
 
     initial_remaining = pool.remaining
 
-    # Draw all layouts for env 0 (variant 0) and env 1 (variant 1)
-    env_ids = torch.tensor([0, 1] * initial_remaining)
-    pool.sample_without_replacement(len(env_ids), env_ids=env_ids)
+    for _ in range(initial_remaining):
+        pool.sample_without_replacement(2)
 
     # Pool should be exhausted now; request more to trigger refill
-    env_ids_more = torch.tensor([0, 1])
-    draws = pool.sample_without_replacement(2, env_ids=env_ids_more)
+    draws = pool.sample_without_replacement(2)
     assert len(draws) == 2, "Pool should refill and return requested layouts"
 
 
-def test_pooled_placer_homogeneous_unaffected_by_num_envs():
-    """Homogeneous pool should work the same whether num_envs is passed or not."""
+def test_pooled_placer_reusable_layouts_allocate_complete_env_rounds():
+    """Reusable layouts should still expose equal without-replacement capacity per env."""
     desk = _make_desk()
     box = DummyObject(
         name="box",
@@ -363,9 +363,62 @@ def test_pooled_placer_homogeneous_unaffected_by_num_envs():
     placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
 
     pool = PooledObjectPlacer(objects=[desk, box], placer_params=placer_params, pool_size=10, num_envs=4)
-    assert not pool.is_heterogeneous
-    draws = pool.sample_without_replacement(3)
-    assert len(draws) == 3
+    assert pool.remaining == 3
+    draws = pool.sample_without_replacement(4)
+    assert len(draws) == 4
+    assert pool.remaining == 2
+
+
+def test_pooled_placer_reusable_layouts_keep_partial_valid_results():
+    """Reusable layouts should not be dropped when fewer than num_envs are valid."""
+    desk = _make_desk()
+    box = DummyObject(
+        name="box",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)),
+    )
+    box.add_relation(On(desk, clearance_m=0.01))
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
+
+    pool = PooledObjectPlacer(objects=[desk, box], placer_params=placer_params, pool_size=4, num_envs=4)
+    pool._layout_pools = {env_id: [] for env_id in range(4)}
+    pool._layout_cursors = {env_id: 0 for env_id in range(4)}
+
+    layouts = [
+        PlacementResult(success=True, positions={box: (float(i), 0.0, 0.0)}, final_loss=0.0, attempts=1)
+        for i in range(3)
+    ]
+    pool._store_reusable_results(layouts)
+
+    assert sum(len(pool._layout_pools[env_id]) for env_id in range(4)) == 3
+    assert pool.remaining == 0
+
+
+def test_pooled_placer_mixed_heterogeneous_and_homogeneous_objects():
+    """A pool with mixed object types should match only per-env geometry by env."""
+    desk = _make_desk()
+    small = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1))
+    large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.25, 0.25, 0.25))
+    hetero = HeterogeneousDummyObject(name="hetero", bboxes=[small, large])
+    hetero.add_relation(On(desk, clearance_m=0.01))
+
+    box = DummyObject(
+        name="box",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.15)),
+    )
+    box.add_relation(On(desk, clearance_m=0.01))
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
+
+    pool = PooledObjectPlacer(objects=[desk, hetero, box], placer_params=placer_params, pool_size=40, num_envs=4)
+
+    draws = pool.sample_without_replacement(4)
+    assert len(draws) == 4
+    for draw in draws:
+        assert hetero in draw.positions
+        assert box in draw.positions
 
 
 # ---------------------------------------------------------------------------
@@ -396,11 +449,9 @@ def test_pooled_placer_multi_set_different_variant_counts():
     placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
 
     pool = PooledObjectPlacer(objects=[desk, bottles, boxes], placer_params=placer_params, pool_size=50, num_envs=6)
-    assert pool.is_heterogeneous
     assert pool.remaining > 0
 
-    env_ids = torch.tensor([0, 1, 2, 3, 4, 5])
-    draws = pool.sample_without_replacement(6, env_ids=env_ids)
+    draws = pool.sample_without_replacement(6)
     assert len(draws) == 6
     for d in draws:
         assert bottles in d.positions
@@ -456,32 +507,26 @@ def test_pooled_placer_multi_set_refill():
 
     # Drain pool then request more to trigger refill
     initial_remaining = pool.remaining
-    env_ids = torch.tensor(list(range(6)) * initial_remaining)
-    pool.sample_without_replacement(len(env_ids), env_ids=env_ids)
+    for _ in range(initial_remaining):
+        pool.sample_without_replacement(6)
 
-    env_ids_more = torch.tensor([0, 1, 2, 3, 4, 5])
-    draws = pool.sample_without_replacement(6, env_ids=env_ids_more)
+    draws = pool.sample_without_replacement(6)
     assert len(draws) == 6
 
 
-def test_pooled_placer_per_env_pools_isolated():
-    """Each env_id should have its own independent pool of layouts."""
+def test_pooled_placer_per_env_pools_advance_in_complete_rounds():
+    """Every env pool cursor should advance together."""
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
     initial_remaining = pool.remaining
 
-    # Draw only from env 0 and env 1; env 2 and 3 should be unaffected.
-    env_ids = torch.tensor([0, 1])
-    pool.sample_without_replacement(2, env_ids=env_ids)
+    pool.sample_without_replacement(4)
 
-    # `remaining` reports the min across all envs.  Env 0 and 1 each lost one
-    # layout, so the min should have decreased by 1.
+    # One complete round was consumed, so every env lost one layout.
     assert pool.remaining == initial_remaining - 1
 
-    # Drawing from env 2 and 3 should still work from their full pools.
-    env_ids_23 = torch.tensor([2, 3])
-    draws = pool.sample_without_replacement(2, env_ids=env_ids_23)
-    assert len(draws) == 2
+    draws = pool.sample_without_replacement(4)
+    assert len(draws) == 4
     for d in draws:
         assert hetero in d.positions
diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py
index 3ac33d0fb..9e0482ee1 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -233,7 +233,7 @@ def test_solve_and_place_objects_handles_multiple_env_ids():
 
     solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3)
     placer_params = ObjectPlacerParams(solver_params=solver_params)
-    pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=10)
+    pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=12, num_envs=num_envs)
 
     solve_and_place_objects(env, env_ids, objects, pool)
 
@@ -325,7 +325,7 @@ def test_resolve_on_reset_false_applies_pose_per_env():
 
     solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3)
     placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
-    pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=20)
+    pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=21, num_envs=num_envs)
 
     layouts = pool.sample_with_replacement(num_envs)
     assert len(layouts) == num_envs

From 20308dee751449b8a271ca9a54775db4894be59f Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 14:57:04 -0700
Subject: [PATCH 18/46] update comments and names

---
 isaaclab_arena/assets/dummy_object.py         |  2 +-
 isaaclab_arena/assets/object_base.py          |  9 ++---
 isaaclab_arena/assets/object_set.py           | 30 +++++-----------
 isaaclab_arena/relations/object_placer.py     | 35 ++++++++++---------
 .../relations/pooled_object_placer.py         |  4 +--
 .../tests/test_heterogeneous_placement.py     | 13 ++++---
 isaaclab_arena/tests/test_placement_events.py |  2 +-
 7 files changed, 45 insertions(+), 50 deletions(-)

diff --git a/isaaclab_arena/assets/dummy_object.py b/isaaclab_arena/assets/dummy_object.py
index 6551ecd6f..50bd9b939 100644
--- a/isaaclab_arena/assets/dummy_object.py
+++ b/isaaclab_arena/assets/dummy_object.py
@@ -43,7 +43,7 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox:
         return self.bounding_box
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
-        """Get per-environment local bounding boxes (expanded from single bbox)."""
+        """Mirror ObjectBase.get_bounding_box_per_env for this test double."""
         bbox = self.get_bounding_box()
         return AxisAlignedBoundingBox(
             min_point=bbox.min_point.expand(num_envs, 3),
diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py
index 35ecc58b1..65801874a 100644
--- a/isaaclab_arena/assets/object_base.py
+++ b/isaaclab_arena/assets/object_base.py
@@ -74,11 +74,12 @@ def get_world_bounding_box(self) -> AxisAlignedBoundingBox:
         ...
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
-        """Get per-environment local bounding boxes.
+        """Get local bounding boxes for each environment.
 
-        For homogeneous objects the single local bbox is expanded to ``(num_envs, 3)``.
-        ``RigidObjectSet`` overrides this to return the actual bbox of each env's
-        variant, enabling heterogeneous placement.
+        This default implementation is for objects with the same geometry in
+        every environment: it expands the single local bbox to ``(num_envs, 3)``.
+        ``RigidObjectSet`` overrides this to return the bbox for each env's
+        assigned variant.
 
         Args:
             num_envs: Number of environments.
diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 10e05d61e..868160652 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -69,13 +69,13 @@ def __init__(
 
         self.objects: list[Object] = objects
         self.random_choice = random_choice
-        self.heterogeneous_bbox: bool = len(objects) > 1
+        self.has_env_specific_bboxes: bool = len(objects) > 1
 
-        if self.heterogeneous_bbox and self.random_choice:
+        if self.has_env_specific_bboxes and self.random_choice:
             raise ValueError(
                 f"RigidObjectSet '{name}': random_choice=True is not supported with heterogeneous "
                 "placement (len(objects) > 1). The placement pool assumes round-robin variant "
-                "assignment (env_idx % num_variants) which conflicts with random spawning order."
+                "assignment (env_index % num_variants) which conflicts with random spawning order."
             )
 
         # Set default prim_path if not provided
@@ -100,34 +100,22 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """
         return max(self.objects, key=lambda obj: obj.get_bounding_box().size[0, 2].item()).get_bounding_box()
 
-    def get_variant_indices(self, num_envs: int, seed: int | None = None) -> list[int]:
+    def get_variant_indices(self, num_envs: int) -> list[int]:
         """Return which member object index is assigned to each environment.
 
-        When ``random_choice`` is False the mapping is round-robin
-        (``env_idx % len(objects)``).  When True, a random permutation is
-        generated (and cached so repeated calls with the same ``num_envs``
-        are deterministic within a session).
+        Multi-variant sets use round-robin assignment
+        (``env_index % len(objects)``). ``random_choice=True`` with multiple
+        variants is rejected in ``__init__`` because placement needs to know
+        the assigned variant for each env.
 
         Args:
             num_envs: Number of environments.
-            seed: Optional RNG seed for reproducible variant assignment
-                when ``random_choice`` is True. If None, uses the global
-                torch RNG.
 
         Returns:
             List of length ``num_envs`` with indices into ``self.objects``.
         """
         n = len(self.objects)
-        if not self.random_choice:
-            return [i % n for i in range(num_envs)]
-
-        if not hasattr(self, "_cached_variant_indices") or len(self._cached_variant_indices) != num_envs:
-            generator = None
-            if seed is not None:
-                generator = torch.Generator()
-                generator.manual_seed(seed)
-            self._cached_variant_indices = [int(i) for i in torch.randint(n, (num_envs,), generator=generator).tolist()]
-        return self._cached_variant_indices
+        return [i % n for i in range(num_envs)]
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         """Get the actual bounding box for each env's variant.
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index 20d23c849..b36ca10a2 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -27,14 +27,14 @@
     from isaaclab_arena.assets.object_base import ObjectBase
 
 
-def _is_heterogeneous(obj: ObjectBase) -> bool:
+def _has_env_specific_bboxes(obj: ObjectBase) -> bool:
     """Return True if *obj* provides per-env variant geometry.
 
-    ``RigidObjectSet`` (and test doubles) set ``heterogeneous_bbox = True``
+    ``RigidObjectSet`` (and test doubles) set ``has_env_specific_bboxes = True``
     to signal that ``get_bounding_box_per_env`` returns different bboxes
     across environments.
     """
-    return getattr(obj, "heterogeneous_bbox", False)
+    return getattr(obj, "has_env_specific_bboxes", False)
 
 
 @dataclass
@@ -138,10 +138,10 @@ def place(
         max_attempts = self.params.max_placement_attempts
         num_candidates = max_attempts * num_results
 
-        # Detect heterogeneous objects (e.g. RigidObjectSet with per-env variants).
-        heterogeneous = result_per_env and any(_is_heterogeneous(obj) for obj in objects)
+        # Detect objects such as RigidObjectSet that expose different bboxes per env.
+        uses_env_specific_bboxes = result_per_env and any(_has_env_specific_bboxes(obj) for obj in objects)
 
-        if heterogeneous:
+        if uses_env_specific_bboxes:
             results_per_env = self._place_heterogeneous(
                 objects,
                 anchor_objects_set,
@@ -222,8 +222,9 @@ def _place_heterogeneous(
     ) -> list[PlacementResult]:
         """Per-env placement: each candidate is tied to its env's object variants.
 
-        Batch layout: candidates [e * max_attempts : (e+1) * max_attempts] belong
-        to env *e*. Per-row bboxes reflect each env's actual variant geometry.
+        Batch layout: candidates
+        ``[cur_env * max_attempts : (cur_env + 1) * max_attempts]`` belong
+        to ``cur_env``. Per-row bboxes reflect each env's actual variant geometry.
 
         Args:
             env_bboxes: When provided, uses these bboxes directly instead of
@@ -234,10 +235,10 @@ def _place_heterogeneous(
             env_bboxes = {obj: obj.get_bounding_box_per_env(num_envs) for obj in objects}
 
         # Expand into per-candidate bboxes (num_candidates, 3): repeat each env's
-        # bbox max_attempts times so rows [e*A:(e+1)*A] share env e's geometry.
+        # bbox max_attempts times so candidates for that env share its geometry.
         candidate_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {}
         for obj, bbox in env_bboxes.items():
-            # bbox.min_point is (num_envs, 3) → repeat_interleave → (num_candidates, 3)
+            # bbox.min_point is (num_envs, 3) -> repeat_interleave -> (num_candidates, 3)
             min_pt = bbox.min_point.repeat_interleave(max_attempts, dim=0)
             max_pt = bbox.max_point.repeat_interleave(max_attempts, dim=0)
             candidate_bboxes[obj] = AxisAlignedBoundingBox(min_point=min_pt, max_point=max_pt)
@@ -245,14 +246,14 @@ def _place_heterogeneous(
         # Generate initial positions; each candidate uses its env's bbox.
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
         for candidate_idx in range(num_candidates):
-            env_idx = candidate_idx // max_attempts
+            cur_env = candidate_idx // max_attempts
             if generator is not None:
                 generator.manual_seed(self.params.placement_seed + candidate_idx)
             # Slice single-env bboxes for this candidate's env.
             env_child_bboxes = {
                 obj: AxisAlignedBoundingBox(
-                    min_point=env_bboxes[obj].min_point[env_idx : env_idx + 1],
-                    max_point=env_bboxes[obj].max_point[env_idx : env_idx + 1],
+                    min_point=env_bboxes[obj].min_point[cur_env : cur_env + 1],
+                    max_point=env_bboxes[obj].max_point[cur_env : cur_env + 1],
                 )
                 for obj in objects
             }
@@ -266,13 +267,13 @@ def _place_heterogeneous(
 
         # Select best candidate per env.
         results: list[PlacementResult] = []
-        for env_idx in range(num_envs):
-            start = env_idx * max_attempts
+        for cur_env in range(num_envs):
+            start = cur_env * max_attempts
             # Slice single-env bboxes for validation of this env's candidates.
             env_bbox_overrides = {
                 obj: AxisAlignedBoundingBox(
-                    min_point=env_bboxes[obj].min_point[env_idx : env_idx + 1],
-                    max_point=env_bboxes[obj].max_point[env_idx : env_idx + 1],
+                    min_point=env_bboxes[obj].min_point[cur_env : cur_env + 1],
+                    max_point=env_bboxes[obj].max_point[cur_env : cur_env + 1],
                 )
                 for obj in objects
             }
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 21c434362..a92604df0 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -54,7 +54,7 @@ def __init__(
         self._objects = list(objects)
         self._placer = ObjectPlacer(params=placer_params)
         self._pool_size = pool_size
-        self._uses_env_specific_bboxes = any(getattr(obj, "heterogeneous_bbox", False) for obj in objects)
+        self._uses_env_specific_bboxes = any(getattr(obj, "has_env_specific_bboxes", False) for obj in objects)
 
         self._num_envs = num_envs if num_envs is not None else 1
         if self._num_envs < 1:
@@ -204,7 +204,7 @@ def _store_env_matched_results(self, all_results: list[PlacementResult], layouts
         if total_valid < total_solved:
             failed_envs = [cur_env for cur_env in self._layout_pools if not self._layout_pools[cur_env]]
             msg = (
-                f"Placement pool (env-matched layouts): solved {total_solved} candidates,"
+                f"Placement pool (env-specific bbox layouts): solved {total_solved} candidates,"
                 f" {total_valid} valid, {total_solved - total_valid} failed validation"
             )
             if failed_envs:
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 249024d8e..1d84f1349 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -3,6 +3,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+# pyright: reportArgumentType=false, reportPrivateUsage=false
+# DummyObject is a lightweight test double for ObjectBase; a few pool tests also
+# inspect internals directly to cover allocation edge cases.
+
 """Tests for heterogeneous object placement with per-env bounding boxes."""
 
 import torch
@@ -35,10 +39,11 @@ class HeterogeneousDummyObject(DummyObject):
     def __init__(self, name: str, bboxes: list[AxisAlignedBoundingBox], **kwargs):
         super().__init__(name=name, bounding_box=bboxes[0], **kwargs)
         self._per_env_bboxes = bboxes
-        self.heterogeneous_bbox = True
+        self.has_env_specific_bboxes = True
         self.objects = bboxes
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
+        """Return env-specific bbox variants for this test double."""
         n_variants = len(self._per_env_bboxes)
         indices = [i % n_variants for i in range(num_envs)]
         min_pts = torch.stack([self._per_env_bboxes[idx].min_point[0] for idx in indices])
@@ -57,7 +62,7 @@ def _make_desk() -> DummyObject:
 
 
 # ---------------------------------------------------------------------------
-# ObjectBase.get_bounding_box_per_env
+# get_bounding_box_per_env default behavior
 # ---------------------------------------------------------------------------
 
 
@@ -163,9 +168,9 @@ def test_placer_heterogeneous_z_height_matches_variant():
 
     desk = _make_desk()
 
-    # "tall" variant: height 0.4 → bottom at z ≈ 0.11 (desk top 0.1 + clearance 0.01)
+    # "tall" variant: height 0.4 -> bottom at z ~0.11 (desk top 0.1 + clearance 0.01)
     tall = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.4))
-    # "short" variant: height 0.1 → bottom at z ≈ 0.11 (same clearance)
+    # "short" variant: height 0.1 -> bottom at z ~0.11 (same clearance)
     short = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.1))
     hetero = HeterogeneousDummyObject(name="hetero", bboxes=[tall, short])
     hetero.add_relation(On(desk, clearance_m=0.01))
diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py
index 9e0482ee1..5844683c9 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -91,7 +91,7 @@ def test_placement_without_seed_multi_env_gives_different_layouts():
     result = placer.place([desk, box1, box2], num_envs=num_envs)
 
     assert isinstance(result, MultiEnvPlacementResult)
-    positions_box1 = [result.results[e].positions[box1] for e in range(num_envs)]
+    positions_box1 = [result.results[env_idx].positions[box1] for env_idx in range(num_envs)]
     any_different = any(positions_box1[i] != positions_box1[j] for i in range(num_envs) for j in range(i + 1, num_envs))
     assert any_different, "Unseeded multi-env placement should produce different positions across environments"
 

From f97bf1deded2dd71d82d1484e5b881c3b1694802 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 15:23:12 -0700
Subject: [PATCH 19/46] clean object names

---
 ..._table_multi_object_no_collision_environment.py | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
index c324cc7e3..751eeceba 100644
--- a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
+++ b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
@@ -63,11 +63,7 @@
     ],
     "cans": [
         "alphabet_soup_can_hope_robolab",
-        "canned_peaches_hope_robolab",
         "corn_can_hope_robolab",
-        "tomato_sauce_can_hope_robolab",
-        "pineapple_slices_can_hope_robolab",
-        "green_beans_can_hope_robolab",
     ],
     "tools": [
         "spoon_handal_robolab",
@@ -75,11 +71,11 @@
         "spoon_2_handal_robolab",
         "measuring_spoon_handal_robolab",
     ],
-    "boxes": [
-        "butter_hope_robolab",
-        "raisin_box_hope_robolab",
-        "yogurt_cup_hope_robolab",
-        "oatmeal_raisin_cookies_hope_robolab",
+    "fruits": [
+        "avocado01_fruits_veggies_robolab",
+        "lemon_01_fruits_veggies_robolab",
+        "orange_01_fruits_veggies_robolab",
+        "pomegranate01_fruits_veggies_robolab",
     ],
 }
 

From 2e816e8c6708ab3a6be4c9d2fb4657676510919e Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 17:11:21 -0700
Subject: [PATCH 20/46] change objectset behavior

---
 isaaclab_arena/assets/object_set.py           | 79 ++++++++++++++-----
 isaaclab_arena/tests/test_object_set.py       | 53 +++++++++++++
 ...e_multi_object_no_collision_environment.py |  2 +-
 3 files changed, 115 insertions(+), 19 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 868160652..90d2d2802 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -29,7 +29,8 @@ def __init__(
         objects: list[Object],
         prim_path: str | None = None,
         scale: tuple[float, float, float] = (1.0, 1.0, 1.0),
-        random_choice: bool = False,
+        random_choice: bool = True,
+        variant_indices_by_env: list[int] | None = None,
         initial_pose: Pose | None = None,
         **kwargs,
     ):
@@ -42,7 +43,9 @@ def __init__(
             scale: The scale of the object set. Note all objects can only have the same scale, if
                 different scales are needed, considering scaling the object USD file.
             random_choice: Whether to randomly choose an object from the object set to spawn in
-                each environment. If False, object is spawned based on the order of objects in the list.
+                each environment. The assignment is sampled once when ``num_envs`` is known and
+                then reused across resets.
+            variant_indices_by_env: Optional fixed variant index for each environment.
             initial_pose: The initial pose of the object from this object set.
         """
         if not self._are_all_objects_type_rigid(objects):
@@ -65,18 +68,19 @@ def __init__(
             self.object_usd_paths = self._modify_assets(objects)
             print(f"Modified object USD paths: {self.object_usd_paths}")
         else:
-            self.object_usd_paths = [object.usd_path for object in objects]
+            self.object_usd_paths = []
+            for obj in objects:
+                assert obj.usd_path is not None
+                self.object_usd_paths.append(obj.usd_path)
 
         self.objects: list[Object] = objects
+        self._member_object_usd_paths: list[str] = list(self.object_usd_paths)
         self.random_choice = random_choice
         self.has_env_specific_bboxes: bool = len(objects) > 1
+        self.variant_indices_by_env: list[int] | None = None
 
-        if self.has_env_specific_bboxes and self.random_choice:
-            raise ValueError(
-                f"RigidObjectSet '{name}': random_choice=True is not supported with heterogeneous "
-                "placement (len(objects) > 1). The placement pool assumes round-robin variant "
-                "assignment (env_index % num_variants) which conflicts with random spawning order."
-            )
+        if variant_indices_by_env is not None:
+            self._set_variant_indices_by_env(variant_indices_by_env)
 
         # Set default prim_path if not provided
         if prim_path is None:
@@ -103,10 +107,9 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox:
     def get_variant_indices(self, num_envs: int) -> list[int]:
         """Return which member object index is assigned to each environment.
 
-        Multi-variant sets use round-robin assignment
-        (``env_index % len(objects)``). ``random_choice=True`` with multiple
-        variants is rejected in ``__init__`` because placement needs to know
-        the assigned variant for each env.
+        Multi-variant sets use one fixed assignment for the lifetime of the
+        object set. By default, each env independently samples one variant
+        once, then keeps it across resets.
 
         Args:
             num_envs: Number of environments.
@@ -114,8 +117,15 @@ def get_variant_indices(self, num_envs: int) -> list[int]:
         Returns:
             List of length ``num_envs`` with indices into ``self.objects``.
         """
-        n = len(self.objects)
-        return [i % n for i in range(num_envs)]
+        if self.variant_indices_by_env is None:
+            self._set_variant_indices_by_env(self._generate_variant_indices(num_envs))
+        elif len(self.variant_indices_by_env) != num_envs:
+            raise ValueError(
+                f"RigidObjectSet '{self.name}' has variant assignments for "
+                f"{len(self.variant_indices_by_env)} envs, got request for {num_envs}."
+            )
+        assert self.variant_indices_by_env is not None
+        return self.variant_indices_by_env
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         """Get the actual bounding box for each env's variant.
@@ -143,10 +153,41 @@ def get_contact_sensor_cfg(self, contact_against_object: ObjectBase | None = Non
         # and we can use the first USD path to find the shallowest rigid body.
         return super().get_contact_sensor_cfg(contact_against_object, usd_path=self.object_usd_paths[0])
 
-    def _are_all_objects_type_rigid(self, objects: list[ObjectBase]) -> bool:
+    def _generate_variant_indices(self, num_envs: int) -> list[int]:
+        n = len(self.objects)
+        if n == 1:
+            return [0 for _ in range(num_envs)]
+        if not self.random_choice:
+            raise ValueError(
+                f"RigidObjectSet '{self.name}' has {n} variants and random_choice=False, "
+                "but no variant_indices_by_env were provided."
+            )
+        return torch.randint(low=0, high=n, size=(num_envs,)).tolist()
+
+    def _set_variant_indices_by_env(self, variant_indices_by_env: list[int]) -> None:
+        n = len(self.objects)
+        if any(idx < 0 or idx >= n for idx in variant_indices_by_env):
+            raise ValueError(
+                f"RigidObjectSet '{self.name}' variant indices must be in [0, {n}); "
+                f"got {variant_indices_by_env}."
+            )
+
+        self.variant_indices_by_env = list(variant_indices_by_env)
+        if self.has_env_specific_bboxes:
+            self.object_usd_paths = [self._member_object_usd_paths[idx] for idx in self.variant_indices_by_env]
+            spawn_cfg = self.object_cfg.spawn if getattr(self, "object_cfg", None) is not None else None
+            if isinstance(spawn_cfg, sim_utils.MultiUsdFileCfg):
+                spawn_cfg.usd_path = self.object_usd_paths
+                spawn_cfg.random_choice = False
+
+    def _are_all_objects_type_rigid(self, objects: list[Object]) -> bool:
         if objects is None or len(objects) == 0:
             raise ValueError(f"Object set {self.name} must contain at least 1 object.")
-        return all(detect_object_type(usd_path=object.usd_path) == ObjectType.RIGID for object in objects)
+        for obj in objects:
+            assert obj.usd_path is not None
+            if detect_object_type(usd_path=obj.usd_path) != ObjectType.RIGID:
+                return False
+        return True
 
     def _generate_rigid_cfg(self) -> RigidObjectCfg:
         assert self.object_type == ObjectType.RIGID
@@ -154,11 +195,12 @@ def _generate_rigid_cfg(self) -> RigidObjectCfg:
             prim_path=self.prim_path,
             spawn=sim_utils.MultiUsdFileCfg(
                 usd_path=self.object_usd_paths,
-                random_choice=self.random_choice,
+                random_choice=self.random_choice if self.variant_indices_by_env is None else False,
                 activate_contact_sensors=True,
             ),
         )
         object_cfg = self._add_initial_pose_to_cfg(object_cfg)
+        assert isinstance(object_cfg, RigidObjectCfg)
         return object_cfg
 
     def _generate_articulation_cfg(self):
@@ -191,6 +233,7 @@ def _asset_modification_possible(self, objects: list[Object]) -> bool:
     def _get_all_rigid_body_depths(self, objects: list[Object]) -> list[int]:
         depths = []
         for asset in objects:
+            assert asset.usd_path is not None
             shallowest_rigid_body = find_shallowest_rigid_body(asset.usd_path)
             depth = shallowest_rigid_body.count("/") - 1 if shallowest_rigid_body else -1
             depths.append(depth)
diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py
index e0dc7b973..15c19d141 100644
--- a/isaaclab_arena/tests/test_object_set.py
+++ b/isaaclab_arena/tests/test_object_set.py
@@ -5,6 +5,7 @@
 
 import os
 import traceback
+from unittest.mock import patch
 
 from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function
 
@@ -16,6 +17,50 @@
 OBJECT_SET_BOTTLES_PRIM_PATH = "/World/envs/env_.*/ObjectSet_Bottles"
 
 
+def _make_object_set_variants():
+    from isaaclab_arena.assets.object import Object
+    from isaaclab_arena.assets.object_base import ObjectType
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+    can_a = Object(name="can_a", object_type=ObjectType.RIGID, usd_path="/tmp/can_a.usd")
+    can_b = Object(name="can_b", object_type=ObjectType.RIGID, usd_path="/tmp/can_b.usd")
+    bbox_a = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.2))
+    bbox_b = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.3))
+    can_a.bounding_box = bbox_a
+    can_b.bounding_box = bbox_b
+    return can_a, can_b, bbox_a, bbox_b
+
+
+def _test_object_set_samples_and_stores_variant_indices(simulation_app):
+    """Variant assignment should be sampled once and reused for spawning and bboxes."""
+    import torch
+
+    from isaaclab_arena.assets.object_base import ObjectType
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+
+    can_a, can_b, bbox_a, bbox_b = _make_object_set_variants()
+    assigned_variant_indices = [1, 0, 1, 1]
+
+    with (
+        patch("isaaclab_arena.assets.object_set.detect_object_type", return_value=ObjectType.RIGID),
+        patch("isaaclab_arena.assets.object_set.find_shallowest_rigid_body", return_value="/rigid"),
+        patch("isaaclab_arena.assets.object_set.torch.randint", return_value=torch.tensor(assigned_variant_indices)),
+    ):
+        obj_set = RigidObjectSet(name="cans", objects=[can_a, can_b])
+        assert obj_set.variant_indices_by_env is None
+        assert obj_set.get_variant_indices(num_envs=4) == assigned_variant_indices
+
+    assert obj_set.object_usd_paths == [can_b.usd_path, can_a.usd_path, can_b.usd_path, can_b.usd_path]
+    spawn_cfg = obj_set.object_cfg.spawn
+    assert getattr(spawn_cfg, "usd_path") == obj_set.object_usd_paths
+    assert getattr(spawn_cfg, "random_choice") is False
+
+    per_env_bbox = obj_set.get_bounding_box_per_env(num_envs=4)
+    assert torch.allclose(per_env_bbox.max_point[0], bbox_b.max_point[0])
+    assert torch.allclose(per_env_bbox.max_point[1], bbox_a.max_point[0])
+    return True
+
+
 def _build_and_reset_env(simulation_app, scene_assets, env_name="object_set_test", task=None):
     """Build arena env with given scene and optional task, then reset. Returns env (caller must close)."""
     from isaaclab_arena.assets.registries import AssetRegistry
@@ -360,6 +405,14 @@ def test_empty_object_set():
     assert result, f"Test {_test_empty_object_set.__name__} failed"
 
 
+def test_object_set_samples_and_stores_variant_indices():
+    result = run_simulation_app_function(
+        _test_object_set_samples_and_stores_variant_indices,
+        headless=HEADLESS,
+    )
+    assert result, f"Test {_test_object_set_samples_and_stores_variant_indices.__name__} failed"
+
+
 def test_articulation_object_set():
     result = run_simulation_app_function(
         _test_articulation_object_set,
diff --git a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
index 751eeceba..466c71b7e 100644
--- a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
+++ b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
@@ -225,7 +225,7 @@ def _build_heterogeneous_objects(self, tabletop_reference, object_names=None):
 
             for set_name, variant_names in HETERO_VARIANT_SETS.items():
                 members = [self.asset_registry.get_asset_by_name(n)() for n in variant_names]
-                obj_set = RigidObjectSet(name=set_name, objects=members)
+                obj_set = RigidObjectSet(name=set_name, objects=members, random_choice=True)
                 obj_set.add_relation(On(tabletop_reference, clearance_m=0.01))
                 placeable_assets.append(obj_set)
 

From f46bfc602aa86c131c99619ca5386f9e35808df0 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 21:19:33 -0700
Subject: [PATCH 21/46] pre-commit fix

---
 isaaclab_arena/assets/object_set.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 90d2d2802..3ccb0827a 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -168,8 +168,7 @@ def _set_variant_indices_by_env(self, variant_indices_by_env: list[int]) -> None
         n = len(self.objects)
         if any(idx < 0 or idx >= n for idx in variant_indices_by_env):
             raise ValueError(
-                f"RigidObjectSet '{self.name}' variant indices must be in [0, {n}); "
-                f"got {variant_indices_by_env}."
+                f"RigidObjectSet '{self.name}' variant indices must be in [0, {n}); got {variant_indices_by_env}."
             )
 
         self.variant_indices_by_env = list(variant_indices_by_env)

From a9bfdde9cd70eaa94796af1f4977f156bae013c1 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 12 May 2026 13:50:48 -0700
Subject: [PATCH 22/46] increase readability

---
 isaaclab_arena/relations/object_placer.py     | 139 ++++++------
 .../relations/pooled_object_placer.py         | 200 ++++++++++--------
 2 files changed, 185 insertions(+), 154 deletions(-)

diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index b36ca10a2..9995083a0 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -27,16 +27,6 @@
     from isaaclab_arena.assets.object_base import ObjectBase
 
 
-def _has_env_specific_bboxes(obj: ObjectBase) -> bool:
-    """Return True if *obj* provides per-env variant geometry.
-
-    ``RigidObjectSet`` (and test doubles) set ``has_env_specific_bboxes = True``
-    to signal that ``get_bounding_box_per_env`` returns different bboxes
-    across environments.
-    """
-    return getattr(obj, "has_env_specific_bboxes", False)
-
-
 @dataclass
 class PlacementCandidate:
     """A single solver result, ranked and selected in ObjectPlacer.place()."""
@@ -79,6 +69,21 @@ def __init__(self, params: ObjectPlacerParams | None = None):
         self.params = params or ObjectPlacerParams()
         self._solver = RelationSolver(params=self.params.solver_params)
 
+    @staticmethod
+    def _resolve_bbox(
+        obj: ObjectBase,
+        overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None,
+    ) -> AxisAlignedBoundingBox:
+        """Return *overrides[obj]* if present, otherwise *obj*'s default bbox.
+
+        Heterogeneous placement passes a single-env override bbox; the
+        homogeneous path and unit tests pass ``None`` and rely on the
+        object's own bbox.
+        """
+        if overrides is not None and obj in overrides:
+            return overrides[obj]
+        return obj.get_bounding_box()
+
     def place(
         self,
         objects: list[ObjectBase],
@@ -138,8 +143,16 @@ def place(
         max_attempts = self.params.max_placement_attempts
         num_candidates = max_attempts * num_results
 
-        # Detect objects such as RigidObjectSet that expose different bboxes per env.
-        uses_env_specific_bboxes = result_per_env and any(_has_env_specific_bboxes(obj) for obj in objects)
+        # Two solve paths produce a list[PlacementResult] of length num_results, sorted by
+        # (is_valid, loss):
+        #   - homogeneous: any solved layout serves any env; pick the top num_results.
+        #   - heterogeneous: some objects vary per env, so each env owns a fixed slice
+        #     of candidates and we pick the best within that slice.
+        # ``has_env_specific_bboxes`` is duck-typed: declared on RigidObjectSet / DummyObject
+        # but not on the abstract ObjectBase, so we read it via getattr.
+        uses_env_specific_bboxes = result_per_env and any(
+            getattr(obj, "has_env_specific_bboxes", False) for obj in objects
+        )
 
         if uses_env_specific_bboxes:
             results_per_env = self._place_heterogeneous(
@@ -177,7 +190,13 @@ def _place_homogeneous(
         num_candidates: int,
         generator: torch.Generator | None,
     ) -> list[PlacementResult]:
-        """Pool-based placement: any valid solution can serve any environment."""
+        """Original batched placement path.
+
+        Solves ``max_attempts * num_results`` candidates in one batched
+        ``solver.solve()`` call, then returns the best ``num_results`` ranked
+        by (is_valid, loss). All envs share the same geometry, so any solved
+        layout can serve any env.
+        """
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
         for candidate_idx in range(num_candidates):
             if generator is not None:
@@ -189,25 +208,30 @@ def _place_homogeneous(
         all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
 
         all_candidates = [
-            PlacementCandidate(all_losses[i], all_positions[i], self._validate_placement(all_positions[i]))
-            for i in range(num_candidates)
+            PlacementCandidate(all_losses[idx], all_positions[idx], self._validate_placement(all_positions[idx]))
+            for idx in range(num_candidates)
         ]
-        all_candidates.sort(key=lambda c: (not c.is_valid, c.loss))
+        all_candidates.sort(key=lambda candidate: (not candidate.is_valid, candidate.loss))
         selected = all_candidates[:num_results]
 
         if self.params.verbose:
-            total_valid = sum(1 for c in all_candidates if c.is_valid)
-            finite_losses = [c.loss for c in all_candidates if math.isfinite(c.loss)]
+            total_valid = sum(1 for candidate in all_candidates if candidate.is_valid)
+            finite_losses = [candidate.loss for candidate in all_candidates if math.isfinite(candidate.loss)]
             mean_loss = sum(finite_losses) / len(finite_losses) if finite_losses else float("inf")
-            n_valid = sum(1 for c in selected if c.is_valid)
+            n_valid = sum(1 for candidate in selected if candidate.is_valid)
             print(
                 f"Solved {num_candidates} candidates in one batch: mean loss = {mean_loss:.6f},"
                 f" {total_valid} valid, selected best {num_results} ({n_valid} valid)"
             )
 
         return [
-            PlacementResult(success=c.is_valid, positions=c.positions, final_loss=c.loss, attempts=max_attempts)
-            for c in selected
+            PlacementResult(
+                success=candidate.is_valid,
+                positions=candidate.positions,
+                final_loss=candidate.loss,
+                attempts=max_attempts,
+            )
+            for candidate in selected
         ]
 
     def _place_heterogeneous(
@@ -222,41 +246,46 @@ def _place_heterogeneous(
     ) -> list[PlacementResult]:
         """Per-env placement: each candidate is tied to its env's object variants.
 
-        Batch layout: candidates
-        ``[cur_env * max_attempts : (cur_env + 1) * max_attempts]`` belong
-        to ``cur_env``. Per-row bboxes reflect each env's actual variant geometry.
+        Candidates ``[cur_env * max_attempts : (cur_env + 1) * max_attempts]``
+        belong to ``cur_env`` and use that env's variant geometry. We solve all
+        ``num_envs * max_attempts`` candidates in one batched ``solver.solve()``
+        call and pick the best candidate within each env's slice.
 
         Args:
-            env_bboxes: When provided, uses these bboxes directly instead of
-                calling ``get_bounding_box_per_env(num_envs)``.  Each bbox must
-                have shape ``(num_envs, 3)``.
+            env_bboxes: Optional pre-tiled per-env bboxes of shape
+                ``(num_envs, 3)``. When ``None`` we call
+                ``get_bounding_box_per_env(num_envs)`` on each object.
         """
         if env_bboxes is None:
             env_bboxes = {obj: obj.get_bounding_box_per_env(num_envs) for obj in objects}
 
-        # Expand into per-candidate bboxes (num_candidates, 3): repeat each env's
-        # bbox max_attempts times so candidates for that env share its geometry.
+        # Build the per-env (1, 3) bbox views once: used both for On-guided
+        # initial-position sampling and for per-env validation below.
+        per_env_bbox_overrides: list[dict[ObjectBase, AxisAlignedBoundingBox]] = [
+            {
+                obj: AxisAlignedBoundingBox(
+                    min_point=env_bboxes[obj].min_point[cur_env : cur_env + 1],
+                    max_point=env_bboxes[obj].max_point[cur_env : cur_env + 1],
+                )
+                for obj in objects
+            }
+            for cur_env in range(num_envs)
+        ]
+
+        # Per-candidate bboxes (num_candidates, 3): each env's row repeated
+        # max_attempts times so candidates for that env share its geometry.
         candidate_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {}
         for obj, bbox in env_bboxes.items():
-            # bbox.min_point is (num_envs, 3) -> repeat_interleave -> (num_candidates, 3)
             min_pt = bbox.min_point.repeat_interleave(max_attempts, dim=0)
             max_pt = bbox.max_point.repeat_interleave(max_attempts, dim=0)
             candidate_bboxes[obj] = AxisAlignedBoundingBox(min_point=min_pt, max_point=max_pt)
 
-        # Generate initial positions; each candidate uses its env's bbox.
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
         for candidate_idx in range(num_candidates):
             cur_env = candidate_idx // max_attempts
             if generator is not None:
                 generator.manual_seed(self.params.placement_seed + candidate_idx)
-            # Slice single-env bboxes for this candidate's env.
-            env_child_bboxes = {
-                obj: AxisAlignedBoundingBox(
-                    min_point=env_bboxes[obj].min_point[cur_env : cur_env + 1],
-                    max_point=env_bboxes[obj].max_point[cur_env : cur_env + 1],
-                )
-                for obj in objects
-            }
+            env_child_bboxes = per_env_bbox_overrides[cur_env]
             initial_positions.append(
                 self._generate_initial_positions(objects, anchor_objects_set, generator, child_bboxes=env_child_bboxes)
             )
@@ -265,27 +294,19 @@ def _place_heterogeneous(
         assert self._solver.last_loss_per_env is not None
         all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
 
-        # Select best candidate per env.
         results: list[PlacementResult] = []
         for cur_env in range(num_envs):
             start = cur_env * max_attempts
-            # Slice single-env bboxes for validation of this env's candidates.
-            env_bbox_overrides = {
-                obj: AxisAlignedBoundingBox(
-                    min_point=env_bboxes[obj].min_point[cur_env : cur_env + 1],
-                    max_point=env_bboxes[obj].max_point[cur_env : cur_env + 1],
-                )
-                for obj in objects
-            }
+            env_bbox_overrides = per_env_bbox_overrides[cur_env]
             env_candidates = [
                 PlacementCandidate(
-                    all_losses[start + j],
-                    all_positions[start + j],
-                    self._validate_placement(all_positions[start + j], bbox_overrides=env_bbox_overrides),
+                    all_losses[start + idx],
+                    all_positions[start + idx],
+                    self._validate_placement(all_positions[start + idx], bbox_overrides=env_bbox_overrides),
                 )
-                for j in range(max_attempts)
+                for idx in range(max_attempts)
             ]
-            env_candidates.sort(key=lambda c: (not c.is_valid, c.loss))
+            env_candidates.sort(key=lambda candidate: (not candidate.is_valid, candidate.loss))
             best = env_candidates[0]
             results.append(
                 PlacementResult(
@@ -295,7 +316,7 @@ def _place_heterogeneous(
 
         if self.params.verbose:
             n_valid = sum(1 for r in results if r.success)
-            print(f"Heterogeneous placement: {n_valid}/{num_envs} env(s) valid")
+            print(f"Solved {num_candidates} candidates in one batch (heterogeneous): {n_valid}/{num_envs} env(s) valid")
 
         return results
 
@@ -459,10 +480,8 @@ def _validate_on_relations(
                 parent = rel.parent
                 if parent not in positions:
                     continue
-                child_bbox = bbox_overrides[obj] if bbox_overrides and obj in bbox_overrides else obj.get_bounding_box()
-                parent_bbox = (
-                    bbox_overrides[parent] if bbox_overrides and parent in bbox_overrides else parent.get_bounding_box()
-                )
+                child_bbox = self._resolve_bbox(obj, bbox_overrides)
+                parent_bbox = self._resolve_bbox(parent, bbox_overrides)
                 child_world = child_bbox.translated(positions[obj])
                 parent_world = parent_bbox.translated(positions[parent])
                 # 1 & 2: Same as OnLossStrategy X/Y band (child's footprint within parent).
@@ -530,8 +549,8 @@ def _validate_no_overlap(
                 if (id(a), id(b)) in on_pairs:
                     continue
 
-                a_bbox = bbox_overrides[a] if bbox_overrides and a in bbox_overrides else a.get_bounding_box()
-                b_bbox = bbox_overrides[b] if bbox_overrides and b in bbox_overrides else b.get_bounding_box()
+                a_bbox = self._resolve_bbox(a, bbox_overrides)
+                b_bbox = self._resolve_bbox(b, bbox_overrides)
                 a_world = a_bbox.translated(positions[a])
                 b_world = b_bbox.translated(positions[b])
 
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index a92604df0..612663e91 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -12,6 +12,7 @@
 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, PlacementResult
+from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
     from isaaclab_arena.assets.object_base import ObjectBase
@@ -20,18 +21,19 @@
 class PooledObjectPlacer:
     """Object placer that keeps a pool of optimized layouts.
 
-    Wraps :class:`ObjectPlacer` and solves object layouts in batches of
-    ``pool_size``, keeping only those that pass validation. The pool is refilled
-    automatically when consumed layouts run out.
+    Storage: ``num_envs`` independent layout pools, each with its own read
+    cursor (this replaces the single ``_layouts`` list + ``_next_idx`` cursor
+    used before heterogeneous placement). Per-env pools are needed because
+    each layout is solved against a specific env's object geometry; sampling
+    is therefore always in env-index order and ``sample_without_replacement``
+    advances every cursor by the same amount on each call.
 
-    Layouts are always stored in per-env pools.  The public sampling methods
-    expose only the replacement strategy; internally, samples are drawn in
-    env-index order.
+    The pool is refilled automatically when an env's queue runs out.
 
     * :meth:`sample_without_replacement` — returns the next *count* layouts
-      sequentially.  Auto-refills when exhausted.
-    * :meth:`sample_with_replacement` — picks *count* layouts at random
-      (non-consuming).  Used for static initial positions.
+      in env-index order (``count`` must be a multiple of ``num_envs``).
+    * :meth:`sample_with_replacement` — picks *count* layouts at random per
+      env-slot (non-consuming). Used for static initial positions.
 
     Args:
         objects: All objects (including anchors) participating in relation solving.
@@ -48,22 +50,26 @@ def __init__(
         pool_size: int = 100,
         num_envs: int | None = None,
     ) -> None:
+        # 1. Validate params.
         if pool_size < 1:
             raise ValueError(f"pool_size must be >= 1, got {pool_size}")
-
-        self._objects = list(objects)
-        self._placer = ObjectPlacer(params=placer_params)
-        self._pool_size = pool_size
+        # ``has_env_specific_bboxes`` is duck-typed (set on RigidObjectSet / DummyObject
+        # but not declared on the abstract ObjectBase), so read it via getattr.
         self._uses_env_specific_bboxes = any(getattr(obj, "has_env_specific_bboxes", False) for obj in objects)
-
+        if self._uses_env_specific_bboxes:
+            assert num_envs is not None, "num_envs is required when layouts use env-specific object variants."
         self._num_envs = num_envs if num_envs is not None else 1
         if self._num_envs < 1:
             raise ValueError(f"num_envs must be >= 1, got {self._num_envs}")
-        if self._uses_env_specific_bboxes:
-            assert num_envs is not None, "num_envs is required when layouts use env-specific object variants."
+
+        # 2. Configure dependencies and per-env storage.
+        self._objects = list(objects)
+        self._placer = ObjectPlacer(params=placer_params)
+        self._pool_size = pool_size
         self._layout_pools: dict[int, list[PlacementResult]] = {cur_env: [] for cur_env in range(self._num_envs)}
         self._layout_cursors: dict[int, int] = {cur_env: 0 for cur_env in range(self._num_envs)}
 
+        # 3. Solve the initial pool and assert every env has at least one layout.
         self._solve_and_store(pool_size)
         for cur_env, pool in self._layout_pools.items():
             if not pool:
@@ -76,6 +82,10 @@ def __init__(
     # Pool storage internals
     # ------------------------------------------------------------------
 
+    def _available_per_env(self) -> list[int]:
+        """Number of unread layouts in each env's pool (length ``num_envs``)."""
+        return [len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] for cur_env in range(self._num_envs)]
+
     def _discard_consumed_layouts(self) -> None:
         """Drop consumed layouts from every env pool before appending new layouts."""
         for cur_env in self._layout_pools:
@@ -84,17 +94,18 @@ def _discard_consumed_layouts(self) -> None:
             self._layout_cursors[cur_env] = 0
 
     def _solve_and_store(self, num_layouts: int) -> None:
-        """Solve layouts and append complete per-env rounds to the pools."""
+        """Solve layouts in batches until every env has ``target_per_env`` unread layouts.
+
+        Each batch contributes (roughly) one round of layouts per env. The
+        outer loop is bounded by ``max_placement_attempts`` to avoid an
+        unbounded refill in pathological configurations.
+        """
         self._discard_consumed_layouts()
         target_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
         max_solve_batches = max(1, self._placer.params.max_placement_attempts)
 
         for _ in range(max_solve_batches):
-            missing_per_env = [
-                target_per_env - (len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env])
-                for cur_env in range(self._num_envs)
-            ]
-            max_missing = max(missing_per_env)
+            max_missing = target_per_env - min(self._available_per_env())
             if max_missing <= 0:
                 return
 
@@ -106,18 +117,12 @@ def _solve_and_store(self, num_layouts: int) -> None:
                 layouts = self._solve_reusable_layouts(batch_size)
                 self._store_reusable_results(layouts)
 
-            available = [
-                len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] for cur_env in range(self._num_envs)
-            ]
-            if min(available) >= target_per_env:
+            if min(self._available_per_env()) >= target_per_env:
                 return
 
-        available = [
-            len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] for cur_env in range(self._num_envs)
-        ]
         raise RuntimeError(
             f"Placement pool could not fill {target_per_env} layouts per env after "
-            f"{max_solve_batches} solve batches. Available per env: {available}."
+            f"{max_solve_batches} solve batches. Available per env: {self._available_per_env()}."
         )
 
     def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
@@ -147,16 +152,21 @@ def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
         return all_layouts
 
     def _store_reusable_results(self, layouts: list[PlacementResult]) -> None:
-        """Distribute reusable layouts across env pools without dropping valid results."""
+        """Distribute reusable layouts across env pools using greedy shortest-first.
+
+        Layouts produced by ``_solve_reusable_layouts`` are interchangeable
+        across envs, so we place each one into whichever pool currently has
+        the fewest unread layouts. This keeps env pools balanced and lets
+        ``sample_without_replacement`` keep advancing in lockstep.
+        """
         if not layouts:
             return
 
+        available = self._available_per_env()
         for layout in layouts:
-            cur_env = min(
-                range(self._num_envs),
-                key=lambda cur_env: len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env],
-            )
+            cur_env = min(range(self._num_envs), key=available.__getitem__)
             self._layout_pools[cur_env].append(layout)
+            available[cur_env] += 1
 
     def _solve_layouts_with_env_bboxes(self, num_layouts: int) -> tuple[list[PlacementResult], int]:
         """Solve layouts tied to each env's actual object geometry.
@@ -164,21 +174,21 @@ def _solve_layouts_with_env_bboxes(self, num_layouts: int) -> tuple[list[Placeme
         Computes bounding boxes for the real ``num_envs`` once, tiles them
         to ``num_layouts`` entries, and solves everything in **one** batched
         ``place()`` call.  Result ``i`` is mapped back to real env
-        ``i % num_envs`` for pool storage.
+        ``i // layouts_per_env`` for pool storage.
         """
-        from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
-
         layouts_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
         total_layouts = layouts_per_env * self._num_envs
 
-        real_bboxes: dict = {obj: obj.get_bounding_box_per_env(self._num_envs) for obj in self._objects}
+        real_bboxes = {obj: obj.get_bounding_box_per_env(self._num_envs) for obj in self._objects}
 
-        tiled_bboxes: dict = {}
-        for obj, bbox in real_bboxes.items():
-            # (num_envs, 3) -> repeat each env's row layouts_per_env times -> (total_layouts, 3)
-            min_pt = bbox.min_point.repeat_interleave(layouts_per_env, dim=0)
-            max_pt = bbox.max_point.repeat_interleave(layouts_per_env, dim=0)
-            tiled_bboxes[obj] = AxisAlignedBoundingBox(min_point=min_pt, max_point=max_pt)
+        # (num_envs, 3) -> repeat each env's row layouts_per_env times -> (total_layouts, 3).
+        tiled_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {
+            obj: AxisAlignedBoundingBox(
+                min_point=bbox.min_point.repeat_interleave(layouts_per_env, dim=0),
+                max_point=bbox.max_point.repeat_interleave(layouts_per_env, dim=0),
+            )
+            for obj, bbox in real_bboxes.items()
+        }
 
         with torch.inference_mode(False):
             result = self._placer.place(
@@ -192,7 +202,14 @@ def _solve_layouts_with_env_bboxes(self, num_layouts: int) -> tuple[list[Placeme
         return all_results, layouts_per_env
 
     def _store_env_matched_results(self, all_results: list[PlacementResult], layouts_per_env: int) -> None:
-        """Store layouts into the env pools they were solved for."""
+        """Store env-matched results into per-env pools, with a best-loss fallback.
+
+        Two passes:
+        1. Append every successful result to its env's pool.
+        2. For any env whose block produced zero successful results, append
+           the block's best-loss candidate (with a warning).
+        """
+        # Pass 1: store successful layouts.
         total_valid = 0
         for i, r in enumerate(all_results):
             cur_env = i // layouts_per_env
@@ -211,36 +228,35 @@ def _store_env_matched_results(self, all_results: list[PlacementResult], layouts
                 msg += f". Envs with zero valid layouts: {failed_envs}"
             print(msg)
 
+        # Pass 2: best-loss fallback for empty env blocks.
         for cur_env in range(self._num_envs):
-            best: PlacementResult | None = None
-            had_valid = False
             start = cur_env * layouts_per_env
-            end = start + layouts_per_env
-            for r in all_results[start:end]:
-                if r.success:
-                    had_valid = True
-                    continue
-                if best is None or r.final_loss < best.final_loss:
-                    best = r
-            if not had_valid and best is not None:
-                print(
-                    f"Warning: env {cur_env} had too few valid layouts; "
-                    f"accepting best-loss fallback (loss={best.final_loss:.6f})."
-                )
-                self._layout_pools[cur_env].append(best)
+            env_block = all_results[start : start + layouts_per_env]
+            if any(r.success for r in env_block):
+                continue
+            best = min(env_block, key=lambda r: r.final_loss, default=None)
+            if best is None:
+                continue
+            print(
+                f"Warning: env {cur_env} had too few valid layouts; "
+                f"accepting best-loss fallback (loss={best.final_loss:.6f})."
+            )
+            self._layout_pools[cur_env].append(best)
 
     # ------------------------------------------------------------------
     # Public API
     # ------------------------------------------------------------------
 
     def sample_without_replacement(self, count: int) -> list[PlacementResult]:
-        """Return the next *count* layouts sequentially (without replacement).
+        """Return the next *count* layouts in env-index order.
 
-        Auto-refills any env pool that does not have enough layouts ahead of
-        its read cursor.
+        Layouts are returned as ``layouts_per_env`` complete rounds of
+        ``[env_0, env_1, ..., env_{num_envs-1}]``, so ``count`` must be a
+        multiple of ``num_envs``. Each round advances every env's cursor by
+        one. Refills any env pool that is short on layouts before reading.
 
         Args:
-            count: Number of layouts to return.
+            count: Number of layouts to return (multiple of ``num_envs``).
 
         Raises:
             ValueError: If *count* is not a complete env round.
@@ -249,41 +265,35 @@ def sample_without_replacement(self, count: int) -> list[PlacementResult]:
         if count % self._num_envs != 0:
             raise ValueError(f"count must be a multiple of num_envs ({self._num_envs}), got {count}")
 
-        sample_env_order = [i % self._num_envs for i in range(count)]
         layouts_per_env = count // self._num_envs
-        needs_refill = any(
-            len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] < layouts_per_env
-            for cur_env in range(self._num_envs)
-        )
-
-        if needs_refill:
+        if min(self._available_per_env()) < layouts_per_env:
             self._solve_and_store(max(self._pool_size, count))
 
         results: list[PlacementResult] = []
-        for cur_env in sample_env_order:
-            idx = self._layout_cursors[cur_env]
-            if idx >= len(self._layout_pools[cur_env]):
-                raise RuntimeError(
-                    f"Placement pool: env {cur_env} has no more valid layouts. "
-                    "The solver is not producing enough valid placements."
-                )
-            results.append(self._layout_pools[cur_env][idx])
-            self._layout_cursors[cur_env] = idx + 1
-
+        for _ in range(layouts_per_env):
+            for cur_env in range(self._num_envs):
+                idx = self._layout_cursors[cur_env]
+                if idx >= len(self._layout_pools[cur_env]):
+                    raise RuntimeError(
+                        f"Placement pool: env {cur_env} has no more valid layouts. "
+                        "The solver is not producing enough valid placements."
+                    )
+                results.append(self._layout_pools[cur_env][idx])
+                self._layout_cursors[cur_env] = idx + 1
         return results
 
     def sample_with_replacement(self, count: int) -> list[PlacementResult]:
-        """Pick *count* layouts at random with replacement (non-consuming).
-
-        Each returned layout is drawn from the per-env pool corresponding to
-        its position in the requested batch.
+        """Pick *count* layouts at random per env-slot (non-consuming).
 
-        Used by ``resolve_on_reset=False`` to assign initial positions
-        that persist across resets.
+        Slot ``i`` is filled by a random pick from env ``i % num_envs``'s
+        pool, so a length-``count`` request walks env indices in the same
+        round-robin order as :meth:`sample_without_replacement`. Used by
+        ``resolve_on_reset=False`` to assign initial positions that persist
+        across resets.
         """
-        sample_env_order = [i % self._num_envs for i in range(count)]
         results: list[PlacementResult] = []
-        for cur_env in sample_env_order:
+        for i in range(count):
+            cur_env = i % self._num_envs
             pool = self._layout_pools[cur_env]
             assert pool, f"Env {cur_env} has no valid layouts to sample from."
             results.append(random.choice(pool))
@@ -291,12 +301,14 @@ def sample_with_replacement(self, count: int) -> list[PlacementResult]:
 
     @property
     def remaining(self) -> int:
-        """Number of layouts not yet consumed by :meth:`sample_without_replacement`.
+        """Number of complete env rounds available to :meth:`sample_without_replacement`.
 
-        Reports the minimum available count across env pools so every env has
-        the same without-replacement capacity.
+        Returns the minimum unread count across env pools (the previous
+        ``remaining`` was a total across one shared list; under per-env
+        storage a single round consumes one layout from every env, so the
+        minimum is what limits without-replacement capacity).
         """
-        return min(len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] for cur_env in self._layout_pools)
+        return min(self._available_per_env())
 
     @property
     def pool_size(self) -> int:

From 53ba645d785d5acc366f7487007ef745c0ecc3c0 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 12 May 2026 17:18:42 -0700
Subject: [PATCH 23/46] address comments

---
 isaaclab_arena/assets/object_set.py           | 11 ++-
 isaaclab_arena/relations/placement_events.py  | 19 +++--
 .../relations/pooled_object_placer.py         | 71 +++++++++++++++----
 isaaclab_arena/tests/test_object_set.py       | 36 +++++++++-
 isaaclab_arena/tests/test_placement_events.py | 29 +++++++-
 5 files changed, 137 insertions(+), 29 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 3ccb0827a..b2d54ff17 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -29,7 +29,7 @@ def __init__(
         objects: list[Object],
         prim_path: str | None = None,
         scale: tuple[float, float, float] = (1.0, 1.0, 1.0),
-        random_choice: bool = True,
+        random_choice: bool = False,
         variant_indices_by_env: list[int] | None = None,
         initial_pose: Pose | None = None,
         **kwargs,
@@ -43,8 +43,8 @@ def __init__(
             scale: The scale of the object set. Note all objects can only have the same scale, if
                 different scales are needed, considering scaling the object USD file.
             random_choice: Whether to randomly choose an object from the object set to spawn in
-                each environment. The assignment is sampled once when ``num_envs`` is known and
-                then reused across resets.
+                each environment. If False, variants are assigned by repeating
+                the member order across environments.
             variant_indices_by_env: Optional fixed variant index for each environment.
             initial_pose: The initial pose of the object from this object set.
         """
@@ -158,10 +158,7 @@ def _generate_variant_indices(self, num_envs: int) -> list[int]:
         if n == 1:
             return [0 for _ in range(num_envs)]
         if not self.random_choice:
-            raise ValueError(
-                f"RigidObjectSet '{self.name}' has {n} variants and random_choice=False, "
-                "but no variant_indices_by_env were provided."
-            )
+            return [env_idx % n for env_idx in range(num_envs)]
         return torch.randint(low=0, high=n, size=(num_envs,)).tolist()
 
     def _set_variant_indices_by_env(self, variant_indices_by_env: list[int]) -> None:
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index 7eb47044c..f8b343f74 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -34,9 +34,10 @@ def solve_and_place_objects(
 ) -> None:
     """Coordinated reset event that draws layouts from the pool and writes poses.
 
-    Registered as a single ``EventTermCfg(mode="reset")``. Each call advances
-    the placement pool by one full env round, then writes poses only for the
-    environments being reset.
+    Registered as a single ``EventTermCfg(mode="reset")``. Env-specific
+    layouts advance by one full env round so each result still matches its
+    absolute env id. Reusable layouts draw only for the environments being
+    reset.
 
     Args:
         env: The Isaac Lab environment.
@@ -47,15 +48,21 @@ def solve_and_place_objects(
     if env_ids is None or len(env_ids) == 0:
         return
 
-    all_results = placement_pool.sample_without_replacement(env.scene.env_origins.shape[0])
+    reset_env_ids = env_ids.tolist()
+    if placement_pool.requires_env_indexed_layouts:
+        all_results = placement_pool.sample_without_replacement(env.scene.env_origins.shape[0])
+        results_by_env = {cur_env: all_results[cur_env] for cur_env in reset_env_ids}
+    else:
+        reset_results = placement_pool.sample_without_replacement(len(reset_env_ids))
+        results_by_env = dict(zip(reset_env_ids, reset_results))
 
     anchor_objects_set = set(get_anchor_objects(objects))
     rotations = {obj: get_rotation_xyzw(obj) for obj in objects if obj not in anchor_objects_set}
 
     zero_velocity = torch.zeros(1, 6, device=env.device)
-    for cur_env in env_ids.tolist():
+    for cur_env in reset_env_ids:
         env_id_tensor = torch.tensor([cur_env], device=env.device)
-        positions = all_results[cur_env].positions
+        positions = results_by_env[cur_env].positions
         for obj, pos in positions.items():
             if obj in anchor_objects_set:
                 continue
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 612663e91..c1610fbef 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -23,15 +23,14 @@ class PooledObjectPlacer:
 
     Storage: ``num_envs`` independent layout pools, each with its own read
     cursor (this replaces the single ``_layouts`` list + ``_next_idx`` cursor
-    used before heterogeneous placement). Per-env pools are needed because
-    each layout is solved against a specific env's object geometry; sampling
-    is therefore always in env-index order and ``sample_without_replacement``
-    advances every cursor by the same amount on each call.
+    used before heterogeneous placement). Env-specific layouts are solved
+    against a fixed env's object geometry and must be sampled in complete env
+    rounds. Reusable layouts can be consumed one at a time.
 
     The pool is refilled automatically when an env's queue runs out.
 
-    * :meth:`sample_without_replacement` — returns the next *count* layouts
-      in env-index order (``count`` must be a multiple of ``num_envs``).
+    * :meth:`sample_without_replacement` — returns the next *count* layouts.
+      Env-specific layouts require ``count`` to be a multiple of ``num_envs``.
     * :meth:`sample_with_replacement` — picks *count* layouts at random per
       env-slot (non-consuming). Used for static initial positions.
 
@@ -86,6 +85,10 @@ def _available_per_env(self) -> list[int]:
         """Number of unread layouts in each env's pool (length ``num_envs``)."""
         return [len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] for cur_env in range(self._num_envs)]
 
+    def _total_available(self) -> int:
+        """Total unread layouts across all env pools."""
+        return sum(self._available_per_env())
+
     def _discard_consumed_layouts(self) -> None:
         """Drop consumed layouts from every env pool before appending new layouts."""
         for cur_env in self._layout_pools:
@@ -248,20 +251,26 @@ def _store_env_matched_results(self, all_results: list[PlacementResult], layouts
     # ------------------------------------------------------------------
 
     def sample_without_replacement(self, count: int) -> list[PlacementResult]:
-        """Return the next *count* layouts in env-index order.
+        """Return the next *count* layouts.
 
-        Layouts are returned as ``layouts_per_env`` complete rounds of
-        ``[env_0, env_1, ..., env_{num_envs-1}]``, so ``count`` must be a
-        multiple of ``num_envs``. Each round advances every env's cursor by
-        one. Refills any env pool that is short on layouts before reading.
+        Env-specific layouts are returned as complete rounds of
+        ``[env_0, env_1, ..., env_{num_envs-1}]`` so each result still maps
+        to the absolute environment it was solved for. Reusable layouts are
+        interchangeable and consume only ``count`` entries.
 
         Args:
-            count: Number of layouts to return (multiple of ``num_envs``).
+            count: Number of layouts to return.
 
         Raises:
-            ValueError: If *count* is not a complete env round.
+            ValueError: If env-specific layouts are requested without a complete env round.
             RuntimeError: If the pool cannot provide *count* layouts after refilling.
         """
+        if self._uses_env_specific_bboxes:
+            return self._sample_env_indexed_without_replacement(count)
+        return self._sample_reusable_without_replacement(count)
+
+    def _sample_env_indexed_without_replacement(self, count: int) -> list[PlacementResult]:
+        """Consume complete env rounds for layouts tied to absolute env ids."""
         if count % self._num_envs != 0:
             raise ValueError(f"count must be a multiple of num_envs ({self._num_envs}), got {count}")
 
@@ -282,6 +291,37 @@ def sample_without_replacement(self, count: int) -> list[PlacementResult]:
                 self._layout_cursors[cur_env] = idx + 1
         return results
 
+    def _sample_reusable_without_replacement(self, count: int) -> list[PlacementResult]:
+        """Consume exactly ``count`` interchangeable layouts."""
+        if self._total_available() < count:
+            self._solve_and_store(max(self._pool_size, count))
+
+        available = self._available_per_env()
+        if sum(available) < count:
+            raise RuntimeError(
+                f"Placement pool has {sum(available)} reusable layouts but {count} were requested. "
+                "The solver is not producing enough valid placements."
+            )
+
+        results: list[PlacementResult] = []
+        for _ in range(count):
+            cur_env = max(range(self._num_envs), key=available.__getitem__)
+            idx = self._layout_cursors[cur_env]
+            if idx >= len(self._layout_pools[cur_env]):
+                raise RuntimeError(
+                    f"Placement pool: env {cur_env} has no more valid layouts. "
+                    "The solver is not producing enough valid placements."
+                )
+            results.append(self._layout_pools[cur_env][idx])
+            self._layout_cursors[cur_env] = idx + 1
+            available[cur_env] -= 1
+        return results
+
+    @property
+    def requires_env_indexed_layouts(self) -> bool:
+        """Whether sampled layouts must be matched back to absolute env ids."""
+        return self._uses_env_specific_bboxes
+
     def sample_with_replacement(self, count: int) -> list[PlacementResult]:
         """Pick *count* layouts at random per env-slot (non-consuming).
 
@@ -314,3 +354,8 @@ def remaining(self) -> int:
     def pool_size(self) -> int:
         """Number of layouts to solve per batch. When the pool runs low, it will solve at least this number of layouts so future samples can reuse the buffer."""
         return self._pool_size
+
+    @property
+    def total_remaining(self) -> int:
+        """Total unread layouts across all env pools."""
+        return self._total_available()
diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py
index 15c19d141..9a8211b6a 100644
--- a/isaaclab_arena/tests/test_object_set.py
+++ b/isaaclab_arena/tests/test_object_set.py
@@ -46,7 +46,7 @@ def _test_object_set_samples_and_stores_variant_indices(simulation_app):
         patch("isaaclab_arena.assets.object_set.find_shallowest_rigid_body", return_value="/rigid"),
         patch("isaaclab_arena.assets.object_set.torch.randint", return_value=torch.tensor(assigned_variant_indices)),
     ):
-        obj_set = RigidObjectSet(name="cans", objects=[can_a, can_b])
+        obj_set = RigidObjectSet(name="cans", objects=[can_a, can_b], random_choice=True)
         assert obj_set.variant_indices_by_env is None
         assert obj_set.get_variant_indices(num_envs=4) == assigned_variant_indices
 
@@ -61,6 +61,32 @@ def _test_object_set_samples_and_stores_variant_indices(simulation_app):
     return True
 
 
+def _test_object_set_default_variant_indices_follow_member_order(simulation_app):
+    """Default object-set assignment should preserve the old deterministic member order."""
+    import torch
+
+    from isaaclab_arena.assets.object_base import ObjectType
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+
+    can_a, can_b, bbox_a, bbox_b = _make_object_set_variants()
+    with (
+        patch("isaaclab_arena.assets.object_set.detect_object_type", return_value=ObjectType.RIGID),
+        patch("isaaclab_arena.assets.object_set.find_shallowest_rigid_body", return_value="/rigid"),
+    ):
+        obj_set = RigidObjectSet(name="ordered_cans", objects=[can_a, can_b])
+        assert obj_set.get_variant_indices(num_envs=5) == [0, 1, 0, 1, 0]
+
+    assert obj_set.object_usd_paths == [can_a.usd_path, can_b.usd_path, can_a.usd_path, can_b.usd_path, can_a.usd_path]
+    spawn_cfg = obj_set.object_cfg.spawn
+    assert getattr(spawn_cfg, "usd_path") == obj_set.object_usd_paths
+    assert getattr(spawn_cfg, "random_choice") is False
+
+    per_env_bbox = obj_set.get_bounding_box_per_env(num_envs=5)
+    assert torch.allclose(per_env_bbox.max_point[0], bbox_a.max_point[0])
+    assert torch.allclose(per_env_bbox.max_point[1], bbox_b.max_point[0])
+    return True
+
+
 def _build_and_reset_env(simulation_app, scene_assets, env_name="object_set_test", task=None):
     """Build arena env with given scene and optional task, then reset. Returns env (caller must close)."""
     from isaaclab_arena.assets.registries import AssetRegistry
@@ -413,6 +439,14 @@ def test_object_set_samples_and_stores_variant_indices():
     assert result, f"Test {_test_object_set_samples_and_stores_variant_indices.__name__} failed"
 
 
+def test_object_set_default_variant_indices_follow_member_order():
+    result = run_simulation_app_function(
+        _test_object_set_default_variant_indices_follow_member_order,
+        headless=HEADLESS,
+    )
+    assert result, f"Test {_test_object_set_default_variant_indices_follow_member_order.__name__} failed"
+
+
 def test_articulation_object_set():
     result = run_simulation_app_function(
         _test_articulation_object_set,
diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py
index 5844683c9..fe9d71c9f 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -247,8 +247,33 @@ def test_solve_and_place_objects_handles_multiple_env_ids():
         )
 
 
-def test_pooled_object_placer_sample_without_replacement_returns_different_layouts():
-    """sample_without_replacement() should return layouts that are likely different."""
+def test_solve_and_place_objects_partial_reset_reusable_pool_consumes_only_reset_envs():
+    """Reusable layouts should not consume a full env round for a partial reset."""
+
+    from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
+    from isaaclab_arena.relations.placement_events import solve_and_place_objects
+    from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer
+    from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
+
+    desk, box1, box2 = _create_test_objects()
+    objects = [desk, box1, box2]
+    num_envs = 4
+    env_ids = torch.tensor([2])
+
+    env = _make_mock_env(num_envs=num_envs)
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3)
+    placer_params = ObjectPlacerParams(solver_params=solver_params)
+    pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=12, num_envs=num_envs)
+
+    available_before = pool.total_remaining
+    solve_and_place_objects(env, env_ids, objects, pool)
+    available_after = pool.total_remaining
+
+    assert available_before - available_after == len(env_ids)
+
+
+def test_pooled_placer_sample_without_replacement_returns_different_layouts():
+    """sample_without_replacement() should return layouts (likely different across draws)."""
 
     from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
     from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer

From 9c11081fbb3696c5416381348c82870c47db99ef Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 13 May 2026 09:13:05 -0700
Subject: [PATCH 24/46] change comments and naming style

---
 isaaclab_arena/assets/object_set.py               |  5 +++--
 isaaclab_arena/relations/pooled_object_placer.py  |  9 ++++-----
 .../tests/test_heterogeneous_placement.py         | 15 +++++++--------
 3 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index b2d54ff17..90d480728 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -108,8 +108,9 @@ def get_variant_indices(self, num_envs: int) -> list[int]:
         """Return which member object index is assigned to each environment.
 
         Multi-variant sets use one fixed assignment for the lifetime of the
-        object set. By default, each env independently samples one variant
-        once, then keeps it across resets.
+        object set. When ``random_choice`` is True, each env independently
+        samples one variant once. Otherwise, assignments repeat the member
+        order across environments.
 
         Args:
             num_envs: Number of environments.
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index c1610fbef..cb1ec5ef6 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -159,8 +159,8 @@ def _store_reusable_results(self, layouts: list[PlacementResult]) -> None:
 
         Layouts produced by ``_solve_reusable_layouts`` are interchangeable
         across envs, so we place each one into whichever pool currently has
-        the fewest unread layouts. This keeps env pools balanced and lets
-        ``sample_without_replacement`` keep advancing in lockstep.
+        the fewest unread layouts. This keeps reusable capacity balanced
+        across env pools.
         """
         if not layouts:
             return
@@ -326,9 +326,8 @@ def sample_with_replacement(self, count: int) -> list[PlacementResult]:
         """Pick *count* layouts at random per env-slot (non-consuming).
 
         Slot ``i`` is filled by a random pick from env ``i % num_envs``'s
-        pool, so a length-``count`` request walks env indices in the same
-        round-robin order as :meth:`sample_without_replacement`. Used by
-        ``resolve_on_reset=False`` to assign initial positions that persist
+        pool, so a length-``count`` request walks env slots in order. Used
+        by ``resolve_on_reset=False`` to assign initial positions that persist
         across resets.
         """
         results: list[PlacementResult] = []
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 1d84f1349..b07ccc0e1 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -40,7 +40,6 @@ def __init__(self, name: str, bboxes: list[AxisAlignedBoundingBox], **kwargs):
         super().__init__(name=name, bounding_box=bboxes[0], **kwargs)
         self._per_env_bboxes = bboxes
         self.has_env_specific_bboxes = True
-        self.objects = bboxes
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         """Return env-specific bbox variants for this test double."""
@@ -99,7 +98,7 @@ def test_heterogeneous_dummy_returns_different_bboxes():
 # ---------------------------------------------------------------------------
 
 
-def test_solver_accepts_env_bboxes():
+def test_relation_solver_uses_env_bboxes():
     """Solver should accept env_bboxes and produce valid results."""
 
     desk = _make_desk()
@@ -134,7 +133,7 @@ def test_solver_accepts_env_bboxes():
 # ---------------------------------------------------------------------------
 
 
-def test_placer_heterogeneous_produces_per_env_results():
+def test_object_placer_heterogeneous_produces_per_env_results():
     """Placer should detect heterogeneous objects and solve per-env."""
 
     desk = _make_desk()
@@ -163,7 +162,7 @@ def test_placer_heterogeneous_produces_per_env_results():
         assert hetero_box in r.positions
 
 
-def test_placer_heterogeneous_z_height_matches_variant():
+def test_object_placer_heterogeneous_z_height_matches_variant():
     """Objects should be placed at z-height matching their env's variant bbox."""
 
     desk = _make_desk()
@@ -255,7 +254,7 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
         ).item(), f"Env {env_idx}: A and X bboxes overlap at positions A={r.positions[obj_a]}, X={r.positions[obj_x]}"
 
 
-def test_homogeneous_path_unchanged():
+def test_object_placer_homogeneous_path_returns_multi_env_result():
     """When no heterogeneous objects exist, the homogeneous path is used."""
 
     desk = _make_desk()
@@ -340,7 +339,7 @@ def test_pooled_placer_heterogeneous_sample_with_replacement():
     assert pool.remaining == initial_remaining, "sample_with_replacement should not consume layouts"
 
 
-def test_pooled_placer_heterogeneous_refill():
+def test_pooled_placer_heterogeneous_sample_without_replacement_triggers_refill():
     """Exhausting a variant sub-pool should trigger a refill."""
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=4, num_envs=2)
@@ -355,7 +354,7 @@ def test_pooled_placer_heterogeneous_refill():
     assert len(draws) == 2, "Pool should refill and return requested layouts"
 
 
-def test_pooled_placer_reusable_layouts_allocate_complete_env_rounds():
+def test_pooled_placer_reusable_layouts_report_complete_env_rounds():
     """Reusable layouts should still expose equal without-replacement capacity per env."""
     desk = _make_desk()
     box = DummyObject(
@@ -489,7 +488,7 @@ def test_pooled_placer_multi_set_sample_with_replacement():
     assert pool.remaining == initial_remaining
 
 
-def test_pooled_placer_multi_set_refill():
+def test_pooled_placer_multi_set_sample_without_replacement_triggers_refill():
     """Exhausting a per-env pool should trigger refill with multi-set objects."""
     desk = _make_desk()
 

From cb4f0e850beb90d939e8b5b92fd3a65f3a48170f Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 13 May 2026 09:30:19 -0700
Subject: [PATCH 25/46] keep 3D collision check

---
 .../relations/relation_loss_strategies.py     | 45 +++++++------------
 isaaclab_arena/relations/relation_solver.py   | 41 ++++++++++-------
 2 files changed, 41 insertions(+), 45 deletions(-)

diff --git a/isaaclab_arena/relations/relation_loss_strategies.py b/isaaclab_arena/relations/relation_loss_strategies.py
index b348e3c6f..6b0d08700 100644
--- a/isaaclab_arena/relations/relation_loss_strategies.py
+++ b/isaaclab_arena/relations/relation_loss_strategies.py
@@ -18,7 +18,7 @@
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
-    from isaaclab_arena.relations.relations import AtPosition, NextTo, On, PositionLimits, Relation
+    from isaaclab_arena.relations.relations import AtPosition, NextTo, On, PositionLimits, Relation, UnaryRelation
 
 from isaaclab_arena.relations.relations import Side
 
@@ -71,7 +71,7 @@ class UnaryRelationLossStrategy(ABC):
     @abstractmethod
     def compute_loss(
         self,
-        relation: "Relation",
+        relation: "UnaryRelation",
         child_pos: torch.Tensor,
         child_bbox: AxisAlignedBoundingBox,
     ) -> torch.Tensor:
@@ -325,11 +325,8 @@ class NoCollisionLossStrategy:
     Computes loss based on:
     1. X overlap: zero when child and parent are separated along X; else overlap length
     2. Y overlap: zero when separated along Y; else overlap length
-    3. Z overlap: zero when separated along Z; else overlap length (skipped when xy_only=True)
-    4. Loss: slope * overlap product (area when xy_only, volume otherwise)
-
-    When ``xy_only=True``, only XY overlap is used — suitable for objects on the
-    same surface where Z overlap is expected and Z gradients would fight the On constraint.
+    3. Z overlap: zero when separated along Z; else overlap length
+    4. Volume loss: slope * (overlap_x * overlap_y * overlap_z)
 
     This is a standalone strategy (not a RelationLossStrategy) because no-overlap
     is a built-in solver behavior, not a user-specified relation.
@@ -351,7 +348,6 @@ def compute_loss(
         child_pos: torch.Tensor,
         child_bbox: AxisAlignedBoundingBox,
         parent_world_bbox: AxisAlignedBoundingBox,
-        xy_only: bool = False,
     ) -> torch.Tensor:
         """Compute loss for no-overlap constraint.
 
@@ -360,8 +356,6 @@ def compute_loss(
             child_pos: Child object position (N, 3) in world coords.
             child_bbox: Child object local bounding box (N=1).
             parent_world_bbox: Parent bounding box in world coordinates.
-            xy_only: If True, compute 2D (XY) overlap area instead of 3D volume.
-                Use for objects on the same surface where Z overlap is expected.
 
         Returns:
             Loss tensor of shape (N,).
@@ -376,6 +370,8 @@ def compute_loss(
         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
 
         # Child world extents
         child_world_min = child_pos + child_bbox.min_point
@@ -384,18 +380,11 @@ def compute_loss(
         # 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)
         overlap_y = interval_overlap_axis_loss(child_world_min[:, 1], child_world_max[:, 1], parent_y_min, parent_y_max)
+        overlap_z = interval_overlap_axis_loss(child_world_min[:, 2], child_world_max[:, 2], parent_z_min, parent_z_max)
 
-        if xy_only:
-            overlap_area = overlap_x * overlap_y
-            total_loss = self.slope * overlap_area
-        else:
-            parent_z_min = parent_world_bbox.min_point[:, 2] - c
-            parent_z_max = parent_world_bbox.max_point[:, 2] + c
-            overlap_z = interval_overlap_axis_loss(
-                child_world_min[:, 2], child_world_max[:, 2], parent_z_min, parent_z_max
-            )
-            overlap_volume = overlap_x * overlap_y * overlap_z
-            total_loss = self.slope * overlap_volume
+        # 2. Volume loss: slope * product of per-axis overlap lengths (overlap volume when slope 1.0)
+        overlap_volume = overlap_x * overlap_y * overlap_z
+        total_loss = self.slope * overlap_volume
 
         if self.debug and child_pos.shape[0] == 1:
             print(
@@ -408,14 +397,12 @@ def compute_loss(
                 f" {child_world_max[0, 1].item():.4f}], parent_y=[{parent_y_min[0].item():.4f},"
                 f" {parent_y_max[0].item():.4f}])"
             )
-            if not xy_only:
-                print(
-                    f"    [NoCollision] Z: overlap={overlap_z[0].item():.6f}"
-                    f" (child_z=[{child_world_min[0, 2].item():.4f},"
-                    f" {child_world_max[0, 2].item():.4f}], parent_z=[{parent_z_min[0].item():.4f},"
-                    f" {parent_z_max[0].item():.4f}])"
-                )
-            print(f"    [NoCollision] loss={total_loss[0].item():.6f}")
+            print(
+                f"    [NoCollision] Z: overlap={overlap_z[0].item():.6f} (child_z=[{child_world_min[0, 2].item():.4f},"
+                f" {child_world_max[0, 2].item():.4f}], parent_z=[{parent_z_min[0].item():.4f},"
+                f" {parent_z_max[0].item():.4f}])"
+            )
+            print(f"    [NoCollision] volume={overlap_volume[0].item():.6f}, loss={total_loss[0].item():.6f}")
 
         return total_loss.squeeze(0) if single_input else total_loss
 
diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index 4f2620152..254ac827b 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import torch
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from isaaclab_arena.relations.relation_loss_strategies import (
     NoCollisionLossStrategy,
@@ -110,7 +110,8 @@ def _compute_total_loss(
 
                 # Handle unary relations (no parent)
                 if isinstance(relation, UnaryRelation):
-                    loss = strategy.compute_loss(
+                    unary_strategy = cast(UnaryRelationLossStrategy, strategy)
+                    loss = unary_strategy.compute_loss(
                         relation=relation,
                         child_pos=child_pos,
                         child_bbox=child_bbox,
@@ -119,6 +120,7 @@ def _compute_total_loss(
                         _print_unary_relation_debug(obj, relation, child_pos[0], loss.mean())
                 # Handle binary relations (with parent) like On, NextTo
                 elif isinstance(relation, Relation):
+                    relation_strategy = cast(RelationLossStrategy, strategy)
                     parent = relation.parent
                     if parent in state.anchor_objects:
                         parent_world_bbox = parent.get_world_bounding_box().to(device)
@@ -126,7 +128,7 @@ def _compute_total_loss(
                         parent_pos = state.get_position(parent)
                         parent_bbox = self._get_bbox(parent, device, env_bboxes)
                         parent_world_bbox = parent_bbox.translated(parent_pos)
-                    loss = strategy.compute_loss(
+                    loss = relation_strategy.compute_loss(
                         relation=relation,
                         child_pos=child_pos,
                         child_bbox=child_bbox,
@@ -152,15 +154,12 @@ def _compute_no_overlap_loss(
         debug: bool = False,
         env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> torch.Tensor:
-        """Compute pairwise no-overlap loss for non-anchor object pairs only.
+        """Compute pairwise no-overlap loss for all non-anchor objects against all other objects.
 
-        Anchor objects (e.g. the table) are excluded: the On relation already
-        controls Z placement relative to anchors, and adding a 3D clearance
-        envelope around the anchor would fight the On constraint in Z.
-
-        Each unique non-anchor pair is evaluated twice (once per direction) so
-        both objects receive gradient.  Uses xy_only=True because objects on the
-        same surface share Z height.
+        Each unique pair is evaluated twice (once per direction):
+        - Non-anchor vs anchor: gradient flows to the non-anchor only.
+        - Non-anchor vs non-anchor: both objects receive gradient by computing
+          the loss in both directions with the other's position detached.
 
         Args:
             state: Current optimization state with object positions.
@@ -174,14 +173,26 @@ def _compute_no_overlap_loss(
         total_loss = torch.zeros(state.batch_size, device=device, dtype=torch.float32)
 
         non_anchor_objects = state.optimizable_objects
+        anchor_objects = list(state.anchor_objects)
 
         for i, child in enumerate(non_anchor_objects):
             child_pos = state.get_position(child)
             child_bbox = self._get_bbox(child, device, env_bboxes)
 
-            # Against other non-anchors (unique pairs, both directions).
-            # Uses xy_only=True because objects on the same surface share Z height;
-            # 3D overlap would generate Z-gradient fighting the On constraint.
+            # Against all anchors
+            for anchor in anchor_objects:
+                anchor_world_bbox = anchor.get_world_bounding_box().to(device)
+                loss = self._no_collision_strategy.compute_loss(
+                    clearance_m=self.params.clearance_m,
+                    child_pos=child_pos,
+                    child_bbox=child_bbox,
+                    parent_world_bbox=anchor_world_bbox,
+                )
+                if debug:
+                    print(f"  [NoOverlap] {child.name} vs {anchor.name}: loss={loss.mean().item():.6f}")
+                total_loss = total_loss + loss
+
+            # Against other non-anchors (unique pairs, both directions)
             for j in range(i + 1, len(non_anchor_objects)):
                 other = non_anchor_objects[j]
                 other_pos = state.get_position(other)
@@ -194,7 +205,6 @@ def _compute_no_overlap_loss(
                     child_pos=child_pos,
                     child_bbox=child_bbox,
                     parent_world_bbox=other_world_bbox,
-                    xy_only=True,
                 )
 
                 # Reverse: gradient flows to other (object j)
@@ -204,7 +214,6 @@ def _compute_no_overlap_loss(
                     child_pos=other_pos,
                     child_bbox=other_bbox,
                     parent_world_bbox=child_world_bbox,
-                    xy_only=True,
                 )
 
                 if debug:

From 5a73e2d2df8f5e2c834b8499085884ef89d035f7 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 13 May 2026 10:16:40 -0700
Subject: [PATCH 26/46] address feedback from review agent

---
 isaaclab_arena/assets/dummy_object.py         |  5 ++-
 isaaclab_arena/assets/object.py               |  3 +-
 isaaclab_arena/assets/object_base.py          |  1 +
 isaaclab_arena/assets/object_set.py           |  4 +-
 isaaclab_arena/relations/object_placer.py     |  8 ++--
 isaaclab_arena/relations/placement_events.py  |  8 +++-
 .../relations/pooled_object_placer.py         | 40 +++----------------
 isaaclab_arena/relations/relation_solver.py   | 12 +++++-
 .../tests/test_heterogeneous_placement.py     | 14 +++++++
 isaaclab_arena/tests/test_placement_events.py | 40 +++++++++++++++++--
 10 files changed, 85 insertions(+), 50 deletions(-)

diff --git a/isaaclab_arena/assets/dummy_object.py b/isaaclab_arena/assets/dummy_object.py
index 50bd9b939..009788f47 100644
--- a/isaaclab_arena/assets/dummy_object.py
+++ b/isaaclab_arena/assets/dummy_object.py
@@ -19,14 +19,15 @@ def __init__(
         name: str,
         bounding_box: AxisAlignedBoundingBox,
         initial_pose: Pose | None = None,
-        relations: list[RelationBase] = [],
+        relations: list[RelationBase] | None = None,
         **kwargs,
     ):
         self.name = name
         self.initial_pose = initial_pose
         self.bounding_box = bounding_box
+        self.has_env_specific_bboxes = False
         assert self.bounding_box is not None
-        self.relations = []
+        self.relations = list(relations or [])
 
     def add_relation(self, relation: RelationBase) -> None:
         self.relations.append(relation)
diff --git a/isaaclab_arena/assets/object.py b/isaaclab_arena/assets/object.py
index bbc2036b4..6123e9f89 100644
--- a/isaaclab_arena/assets/object.py
+++ b/isaaclab_arena/assets/object.py
@@ -32,7 +32,7 @@ def __init__(
         usd_path: str | None = None,
         scale: tuple[float, float, float] = (1.0, 1.0, 1.0),
         initial_pose: Pose | None = None,
-        relations: list[RelationBase] = [],
+        relations: list[RelationBase] | None = None,
         spawner_cfg: SpawnerCfg | None = None,
         **kwargs,
     ):
@@ -55,6 +55,7 @@ def __init__(
         self.spawner_cfg = spawner_cfg
         self.scale = scale
         self.initial_pose = initial_pose
+        self.relations = list(relations or [])
         self.reset_pose = True
         self.spawn_cfg_addon = spawn_cfg_addon
         self.asset_cfg_addon = asset_cfg_addon
diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py
index 65801874a..de552485d 100644
--- a/isaaclab_arena/assets/object_base.py
+++ b/isaaclab_arena/assets/object_base.py
@@ -54,6 +54,7 @@ def __init__(
         self.object_cfg = None
         self.event_cfg = None
         self.relations: list[RelationBase] = []
+        self.has_env_specific_bboxes: bool = False
 
     def get_initial_pose(self) -> Pose | PoseRange | PosePerEnv | None:
         """Return the current initial pose of this object.
diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 90d480728..dd1773484 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -76,7 +76,6 @@ def __init__(
         self.objects: list[Object] = objects
         self._member_object_usd_paths: list[str] = list(self.object_usd_paths)
         self.random_choice = random_choice
-        self.has_env_specific_bboxes: bool = len(objects) > 1
         self.variant_indices_by_env: list[int] | None = None
 
         if variant_indices_by_env is not None:
@@ -95,6 +94,7 @@ def __init__(
             initial_pose=initial_pose,
             **kwargs,
         )
+        self.has_env_specific_bboxes = len(objects) > 1
 
     def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get the bounding box of the object set.
@@ -170,7 +170,7 @@ def _set_variant_indices_by_env(self, variant_indices_by_env: list[int]) -> None
             )
 
         self.variant_indices_by_env = list(variant_indices_by_env)
-        if self.has_env_specific_bboxes:
+        if len(self.objects) > 1:
             self.object_usd_paths = [self._member_object_usd_paths[idx] for idx in self.variant_indices_by_env]
             spawn_cfg = self.object_cfg.spawn if getattr(self, "object_cfg", None) is not None else None
             if isinstance(spawn_cfg, sim_utils.MultiUsdFileCfg):
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index 9995083a0..a7633a636 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -148,11 +148,7 @@ def place(
         #   - homogeneous: any solved layout serves any env; pick the top num_results.
         #   - heterogeneous: some objects vary per env, so each env owns a fixed slice
         #     of candidates and we pick the best within that slice.
-        # ``has_env_specific_bboxes`` is duck-typed: declared on RigidObjectSet / DummyObject
-        # but not on the abstract ObjectBase, so we read it via getattr.
-        uses_env_specific_bboxes = result_per_env and any(
-            getattr(obj, "has_env_specific_bboxes", False) for obj in objects
-        )
+        uses_env_specific_bboxes = result_per_env and any(obj.has_env_specific_bboxes for obj in objects)
 
         if uses_env_specific_bboxes:
             results_per_env = self._place_heterogeneous(
@@ -200,6 +196,7 @@ def _place_homogeneous(
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
         for candidate_idx in range(num_candidates):
             if generator is not None:
+                assert self.params.placement_seed is not None
                 generator.manual_seed(self.params.placement_seed + candidate_idx)
             initial_positions.append(self._generate_initial_positions(objects, anchor_objects_set, generator))
 
@@ -284,6 +281,7 @@ def _place_heterogeneous(
         for candidate_idx in range(num_candidates):
             cur_env = candidate_idx // max_attempts
             if generator is not None:
+                assert self.params.placement_seed is not None
                 generator.manual_seed(self.params.placement_seed + candidate_idx)
             env_child_bboxes = per_env_bbox_overrides[cur_env]
             initial_positions.append(
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index f8b343f74..7866cddb7 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -62,7 +62,13 @@ def solve_and_place_objects(
     zero_velocity = torch.zeros(1, 6, device=env.device)
     for cur_env in reset_env_ids:
         env_id_tensor = torch.tensor([cur_env], device=env.device)
-        positions = results_by_env[cur_env].positions
+        result = results_by_env[cur_env]
+        if not result.success:
+            raise RuntimeError(
+                f"Placement pool returned an invalid layout for env {cur_env} "
+                f"(loss={result.final_loss:.6f}); refusing to write it to simulation."
+            )
+        positions = result.positions
         for obj, pos in positions.items():
             if obj in anchor_objects_set:
                 continue
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index cb1ec5ef6..27e65803c 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -52,9 +52,7 @@ def __init__(
         # 1. Validate params.
         if pool_size < 1:
             raise ValueError(f"pool_size must be >= 1, got {pool_size}")
-        # ``has_env_specific_bboxes`` is duck-typed (set on RigidObjectSet / DummyObject
-        # but not declared on the abstract ObjectBase), so read it via getattr.
-        self._uses_env_specific_bboxes = any(getattr(obj, "has_env_specific_bboxes", False) for obj in objects)
+        self._uses_env_specific_bboxes = any(obj.has_env_specific_bboxes for obj in objects)
         if self._uses_env_specific_bboxes:
             assert num_envs is not None, "num_envs is required when layouts use env-specific object variants."
         self._num_envs = num_envs if num_envs is not None else 1
@@ -131,9 +129,9 @@ def _solve_and_store(self, num_layouts: int) -> None:
     def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
         """Solve layouts that can be used by any env pool.
 
-        When no layouts pass strict validation, the best-loss layouts are
-        accepted with a warning (matching pre-pool behaviour where validation
-        failures were non-fatal).
+        Invalid candidates are discarded. If the solver cannot produce enough
+        valid layouts after the bounded refill attempts, ``_solve_and_store``
+        raises instead of letting invalid placements reach simulation.
         """
         with torch.inference_mode(False):
             result = self._placer.place(self._objects, num_envs=num_layouts, result_per_env=True)
@@ -148,11 +146,7 @@ def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
                 f" {len(valid_layouts)} valid, {num_layouts - len(valid_layouts)} failed validation"
             )
 
-        if valid_layouts:
-            return valid_layouts
-
-        print("Warning: No candidates passed strict validation. Accepting best-loss layouts as fallback.")
-        return all_layouts
+        return valid_layouts
 
     def _store_reusable_results(self, layouts: list[PlacementResult]) -> None:
         """Distribute reusable layouts across env pools using greedy shortest-first.
@@ -205,14 +199,7 @@ def _solve_layouts_with_env_bboxes(self, num_layouts: int) -> tuple[list[Placeme
         return all_results, layouts_per_env
 
     def _store_env_matched_results(self, all_results: list[PlacementResult], layouts_per_env: int) -> None:
-        """Store env-matched results into per-env pools, with a best-loss fallback.
-
-        Two passes:
-        1. Append every successful result to its env's pool.
-        2. For any env whose block produced zero successful results, append
-           the block's best-loss candidate (with a warning).
-        """
-        # Pass 1: store successful layouts.
+        """Store only successful env-matched results into their corresponding pools."""
         total_valid = 0
         for i, r in enumerate(all_results):
             cur_env = i // layouts_per_env
@@ -231,21 +218,6 @@ def _store_env_matched_results(self, all_results: list[PlacementResult], layouts
                 msg += f". Envs with zero valid layouts: {failed_envs}"
             print(msg)
 
-        # Pass 2: best-loss fallback for empty env blocks.
-        for cur_env in range(self._num_envs):
-            start = cur_env * layouts_per_env
-            env_block = all_results[start : start + layouts_per_env]
-            if any(r.success for r in env_block):
-                continue
-            best = min(env_block, key=lambda r: r.final_loss, default=None)
-            if best is None:
-                continue
-            print(
-                f"Warning: env {cur_env} had too few valid layouts; "
-                f"accepting best-loss fallback (loss={best.final_loss:.6f})."
-            )
-            self._layout_pools[cur_env].append(best)
-
     # ------------------------------------------------------------------
     # Public API
     # ------------------------------------------------------------------
diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index 254ac827b..a70055356 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -15,7 +15,7 @@
 )
 from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
 from isaaclab_arena.relations.relation_solver_state import RelationSolverState
-from isaaclab_arena.relations.relations import Relation, RelationBase, UnaryRelation
+from isaaclab_arena.relations.relations import On, Relation, RelationBase, UnaryRelation
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
@@ -174,6 +174,12 @@ def _compute_no_overlap_loss(
 
         non_anchor_objects = state.optimizable_objects
         anchor_objects = list(state.anchor_objects)
+        on_pairs: set[tuple[int, int]] = set()
+        for obj in [*non_anchor_objects, *anchor_objects]:
+            for rel in obj.get_relations():
+                if isinstance(rel, On):
+                    on_pairs.add((id(obj), id(rel.parent)))
+                    on_pairs.add((id(rel.parent), id(obj)))
 
         for i, child in enumerate(non_anchor_objects):
             child_pos = state.get_position(child)
@@ -181,6 +187,8 @@ def _compute_no_overlap_loss(
 
             # Against all anchors
             for anchor in anchor_objects:
+                if (id(child), id(anchor)) in on_pairs:
+                    continue
                 anchor_world_bbox = anchor.get_world_bounding_box().to(device)
                 loss = self._no_collision_strategy.compute_loss(
                     clearance_m=self.params.clearance_m,
@@ -195,6 +203,8 @@ def _compute_no_overlap_loss(
             # Against other non-anchors (unique pairs, both directions)
             for j in range(i + 1, len(non_anchor_objects)):
                 other = non_anchor_objects[j]
+                if (id(child), id(other)) in on_pairs:
+                    continue
                 other_pos = state.get_position(other)
                 other_bbox = self._get_bbox(other, device, env_bboxes)
 
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index b07ccc0e1..13be8887d 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -93,6 +93,20 @@ def test_heterogeneous_dummy_returns_different_bboxes():
     assert torch.allclose(per_env.max_point[1], torch.tensor([0.3, 0.3, 0.3]))
 
 
+def test_dummy_object_preserves_constructor_relations():
+    """DummyObject should keep relations passed at construction time."""
+
+    anchor_relation = IsAnchor()
+    obj = DummyObject(
+        name="anchor",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.1)),
+        relations=[anchor_relation],
+    )
+
+    assert obj.get_relations() == [anchor_relation]
+    assert obj.has_env_specific_bboxes is False
+
+
 # ---------------------------------------------------------------------------
 # Solver with per-env bboxes
 # ---------------------------------------------------------------------------
diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py
index fe9d71c9f..bcb37455d 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -8,6 +8,8 @@
 import torch
 from unittest.mock import MagicMock
 
+import pytest
+
 
 def _create_test_objects():
     """Create a desk (anchor) with two boxes (On + NextTo)."""
@@ -272,6 +274,36 @@ def test_solve_and_place_objects_partial_reset_reusable_pool_consumes_only_reset
     assert available_before - available_after == len(env_ids)
 
 
+def test_solve_and_place_objects_rejects_invalid_pool_layout():
+    """Invalid pool layouts should not be written to simulation."""
+
+    from isaaclab_arena.relations.placement_events import solve_and_place_objects
+    from isaaclab_arena.relations.placement_result import PlacementResult
+
+    desk, box1, box2 = _create_test_objects()
+    objects = [desk, box1, box2]
+    env = _make_mock_env(num_envs=1)
+
+    class InvalidPool:
+        requires_env_indexed_layouts = False
+
+        def sample_without_replacement(self, count: int) -> list[PlacementResult]:
+            assert count == 1
+            return [
+                PlacementResult(
+                    success=False,
+                    positions={box1: (0.0, 0.0, 0.0), box2: (0.0, 0.0, 0.0)},
+                    final_loss=float("nan"),
+                    attempts=1,
+                )
+            ]
+
+    with pytest.raises(RuntimeError, match="invalid layout"):
+        solve_and_place_objects(env, torch.tensor([0]), objects, InvalidPool())
+
+    assert len(env._assets) == 0
+
+
 def test_pooled_placer_sample_without_replacement_returns_different_layouts():
     """sample_without_replacement() should return layouts (likely different across draws)."""
 
@@ -370,8 +402,8 @@ def test_resolve_on_reset_false_applies_pose_per_env():
             assert p.position_xyz is not None, f"Position should not be None for {obj.name}"
 
 
-def test_pooled_object_placer_fallback_when_no_valid_layouts():
-    """PooledObjectPlacer should fall back to best-loss layouts when none pass validation."""
+def test_pooled_placer_raises_when_no_valid_layouts():
+    """PooledObjectPlacer should fail instead of storing invalid layouts."""
 
     from isaaclab_arena.assets.dummy_object import DummyObject
     from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
@@ -403,5 +435,5 @@ def test_pooled_object_placer_fallback_when_no_valid_layouts():
     solver_params = RelationSolverParams(max_iters=50, convergence_threshold=1e-6)
     placer_params = ObjectPlacerParams(solver_params=solver_params, max_placement_attempts=1)
 
-    pool = PooledObjectPlacer(objects=[desk, big1, big2], placer_params=placer_params, pool_size=5)
-    assert pool.remaining > 0, "Pool should contain fallback layouts even when validation fails"
+    with pytest.raises(RuntimeError, match="could not fill"):
+        PooledObjectPlacer(objects=[desk, big1, big2], placer_params=placer_params, pool_size=5)

From 7968c33b258c408fc670da81318df00fc6447827 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 13 May 2026 10:25:32 -0700
Subject: [PATCH 27/46] address new feedback

---
 isaaclab_arena/assets/object_base.py          |  3 +-
 isaaclab_arena/assets/object_set.py           |  2 +-
 isaaclab_arena/relations/placement_events.py  | 13 ++++---
 .../tests/test_heterogeneous_placement.py     | 17 +++++++++
 .../tests/test_no_collision_loss.py           | 27 ++++++++++++++
 isaaclab_arena/tests/test_placement_events.py | 36 +++++++++++++++++++
 6 files changed, 91 insertions(+), 7 deletions(-)

diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py
index de552485d..5243d059d 100644
--- a/isaaclab_arena/assets/object_base.py
+++ b/isaaclab_arena/assets/object_base.py
@@ -54,7 +54,8 @@ def __init__(
         self.object_cfg = None
         self.event_cfg = None
         self.relations: list[RelationBase] = []
-        self.has_env_specific_bboxes: bool = False
+        if not hasattr(self, "has_env_specific_bboxes"):
+            self.has_env_specific_bboxes: bool = False
 
     def get_initial_pose(self) -> Pose | PoseRange | PosePerEnv | None:
         """Return the current initial pose of this object.
diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index dd1773484..dd053f6d1 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -77,6 +77,7 @@ def __init__(
         self._member_object_usd_paths: list[str] = list(self.object_usd_paths)
         self.random_choice = random_choice
         self.variant_indices_by_env: list[int] | None = None
+        self.has_env_specific_bboxes = len(objects) > 1
 
         if variant_indices_by_env is not None:
             self._set_variant_indices_by_env(variant_indices_by_env)
@@ -94,7 +95,6 @@ def __init__(
             initial_pose=initial_pose,
             **kwargs,
         )
-        self.has_env_specific_bboxes = len(objects) > 1
 
     def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get the bounding box of the object set.
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index 7866cddb7..22301c952 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -56,18 +56,21 @@ def solve_and_place_objects(
         reset_results = placement_pool.sample_without_replacement(len(reset_env_ids))
         results_by_env = dict(zip(reset_env_ids, reset_results))
 
-    anchor_objects_set = set(get_anchor_objects(objects))
-    rotations = {obj: get_rotation_xyzw(obj) for obj in objects if obj not in anchor_objects_set}
-
-    zero_velocity = torch.zeros(1, 6, device=env.device)
     for cur_env in reset_env_ids:
-        env_id_tensor = torch.tensor([cur_env], device=env.device)
         result = results_by_env[cur_env]
         if not result.success:
             raise RuntimeError(
                 f"Placement pool returned an invalid layout for env {cur_env} "
                 f"(loss={result.final_loss:.6f}); refusing to write it to simulation."
             )
+
+    anchor_objects_set = set(get_anchor_objects(objects))
+    rotations = {obj: get_rotation_xyzw(obj) for obj in objects if obj not in anchor_objects_set}
+
+    zero_velocity = torch.zeros(1, 6, device=env.device)
+    for cur_env in reset_env_ids:
+        env_id_tensor = torch.tensor([cur_env], device=env.device)
+        result = results_by_env[cur_env]
         positions = result.positions
         for obj, pos in positions.items():
             if obj in anchor_objects_set:
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 13be8887d..fcbe26964 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -107,6 +107,23 @@ def test_dummy_object_preserves_constructor_relations():
     assert obj.has_env_specific_bboxes is False
 
 
+def test_object_preserves_constructor_relations():
+    """Object should keep relations passed at construction time."""
+    from isaaclab_arena.assets.object import Object
+    from isaaclab_arena.assets.object_base import ObjectType
+
+    anchor_relation = IsAnchor()
+    obj = Object(
+        name="rigid_object",
+        object_type=ObjectType.RIGID,
+        usd_path="/tmp/rigid_object.usd",
+        relations=[anchor_relation],
+    )
+
+    assert obj.get_relations() == [anchor_relation]
+    assert obj.has_env_specific_bboxes is False
+
+
 # ---------------------------------------------------------------------------
 # Solver with per-env bboxes
 # ---------------------------------------------------------------------------
diff --git a/isaaclab_arena/tests/test_no_collision_loss.py b/isaaclab_arena/tests/test_no_collision_loss.py
index 6383a92ee..c7397abdd 100644
--- a/isaaclab_arena/tests/test_no_collision_loss.py
+++ b/isaaclab_arena/tests/test_no_collision_loss.py
@@ -11,6 +11,7 @@
 from isaaclab_arena.relations.relation_loss_strategies import NoCollisionLossStrategy
 from isaaclab_arena.relations.relation_solver import RelationSolver
 from isaaclab_arena.relations.relation_solver_params import RelationSolverParams
+from isaaclab_arena.relations.relation_solver_state import RelationSolverState
 from isaaclab_arena.relations.relations import IsAnchor, On
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 from isaaclab_arena.utils.pose import Pose
@@ -216,6 +217,32 @@ def test_solver_respects_clearance_m():
     ).item(), f"Boxes should be at least 5 cm apart; box_a at {pos_a}, box_b at {pos_b}"
 
 
+def test_no_overlap_skips_direct_on_non_anchor_pair():
+    """No-overlap should not fight an On relation whose parent is also movable."""
+    table = _create_table()
+    table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)))
+    table.add_relation(IsAnchor())
+
+    support = _create_box("support")
+    child = _create_box("child")
+    support.add_relation(On(table, clearance_m=0.0))
+    child.add_relation(On(support, clearance_m=0.0))
+
+    objects = [table, support, child]
+    initial_positions = [{
+        table: (0.0, 0.0, 0.0),
+        support: (2.0, 2.0, 0.0),
+        child: (2.0, 2.0, 0.0),
+    }]
+    state = RelationSolverState(objects, initial_positions, device=torch.device("cpu"))
+    solver = RelationSolver(params=RelationSolverParams(max_iters=0))
+
+    loss = solver._compute_no_overlap_loss(state)
+
+    assert torch.isfinite(loss).all()
+    assert torch.allclose(loss, torch.zeros_like(loss))
+
+
 def test_negative_clearance_m_raises():
     """Negative clearance_m should be rejected."""
     import pytest
diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py
index bcb37455d..505e4c9ec 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -304,6 +304,42 @@ def sample_without_replacement(self, count: int) -> list[PlacementResult]:
     assert len(env._assets) == 0
 
 
+def test_solve_and_place_objects_rejects_invalid_layout_before_partial_write():
+    """Invalid layouts should abort the whole reset before any env is written."""
+
+    from isaaclab_arena.relations.placement_events import solve_and_place_objects
+    from isaaclab_arena.relations.placement_result import PlacementResult
+
+    desk, box1, box2 = _create_test_objects()
+    objects = [desk, box1, box2]
+    env = _make_mock_env(num_envs=2)
+
+    class PartiallyInvalidPool:
+        requires_env_indexed_layouts = False
+
+        def sample_without_replacement(self, count: int) -> list[PlacementResult]:
+            assert count == 2
+            return [
+                PlacementResult(
+                    success=True,
+                    positions={box1: (0.1, 0.1, 0.1), box2: (0.2, 0.2, 0.2)},
+                    final_loss=0.0,
+                    attempts=1,
+                ),
+                PlacementResult(
+                    success=False,
+                    positions={box1: (0.0, 0.0, 0.0), box2: (0.0, 0.0, 0.0)},
+                    final_loss=float("nan"),
+                    attempts=1,
+                ),
+            ]
+
+    with pytest.raises(RuntimeError, match="invalid layout"):
+        solve_and_place_objects(env, torch.tensor([0, 1]), objects, PartiallyInvalidPool())
+
+    assert len(env._assets) == 0
+
+
 def test_pooled_placer_sample_without_replacement_returns_different_layouts():
     """sample_without_replacement() should return layouts (likely different across draws)."""
 

From c11eee33a108e132db12bef605eeae83212b1b9c Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 13 May 2026 10:32:27 -0700
Subject: [PATCH 28/46] add bbox guard

---
 isaaclab_arena/assets/object_base.py                 | 4 ++--
 isaaclab_arena/relations/placement_events.py         | 2 ++
 isaaclab_arena/tests/test_heterogeneous_placement.py | 1 +
 isaaclab_arena/tests/test_no_collision_loss.py       | 3 +++
 4 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py
index 5243d059d..cec3f6ecb 100644
--- a/isaaclab_arena/assets/object_base.py
+++ b/isaaclab_arena/assets/object_base.py
@@ -37,6 +37,8 @@
 class ObjectBase(Asset, ABC):
     """Parent class for (spawnable) Object and ObjectReference."""
 
+    has_env_specific_bboxes: bool = False
+
     def __init__(
         self,
         name: str,
@@ -54,8 +56,6 @@ def __init__(
         self.object_cfg = None
         self.event_cfg = None
         self.relations: list[RelationBase] = []
-        if not hasattr(self, "has_env_specific_bboxes"):
-            self.has_env_specific_bboxes: bool = False
 
     def get_initial_pose(self) -> Pose | PoseRange | PosePerEnv | None:
         """Return the current initial pose of this object.
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index 22301c952..b85b6ae31 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -56,6 +56,8 @@ def solve_and_place_objects(
         reset_results = placement_pool.sample_without_replacement(len(reset_env_ids))
         results_by_env = dict(zip(reset_env_ids, reset_results))
 
+    # PooledObjectPlacer stores only successful layouts; this pre-pass keeps
+    # reset writes atomic if a custom pool or future regression violates that invariant.
     for cur_env in reset_env_ids:
         result = results_by_env[cur_env]
         if not result.success:
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index fcbe26964..82e822bb7 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -116,6 +116,7 @@ def test_object_preserves_constructor_relations():
     obj = Object(
         name="rigid_object",
         object_type=ObjectType.RIGID,
+        # Explicit object_type avoids USD inspection; the path is never opened in this constructor-relations test.
         usd_path="/tmp/rigid_object.usd",
         relations=[anchor_relation],
     )
diff --git a/isaaclab_arena/tests/test_no_collision_loss.py b/isaaclab_arena/tests/test_no_collision_loss.py
index c7397abdd..a52ad09bb 100644
--- a/isaaclab_arena/tests/test_no_collision_loss.py
+++ b/isaaclab_arena/tests/test_no_collision_loss.py
@@ -3,6 +3,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+# pyright: reportPrivateUsage=false
+# This file intentionally probes RelationSolver internals for focused loss checks.
+
 """Tests for NoCollisionLossStrategy and RelationSolver built-in no-overlap behavior."""
 
 import torch

From 712c39af627e91efb903f3066f40d129c5c60bc2 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 13 May 2026 13:15:13 -0700
Subject: [PATCH 29/46] seperate anchor bahavior

---
 isaaclab_arena/relations/placement_events.py  | 10 ----
 .../relations/pooled_object_placer.py         | 40 +++++++++----
 .../tests/test_no_collision_loss.py           |  5 +-
 isaaclab_arena/tests/test_placement_events.py | 59 ++++---------------
 ...e_multi_object_no_collision_environment.py |  2 +
 5 files changed, 42 insertions(+), 74 deletions(-)

diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index b85b6ae31..cf0a85da3 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -56,16 +56,6 @@ def solve_and_place_objects(
         reset_results = placement_pool.sample_without_replacement(len(reset_env_ids))
         results_by_env = dict(zip(reset_env_ids, reset_results))
 
-    # PooledObjectPlacer stores only successful layouts; this pre-pass keeps
-    # reset writes atomic if a custom pool or future regression violates that invariant.
-    for cur_env in reset_env_ids:
-        result = results_by_env[cur_env]
-        if not result.success:
-            raise RuntimeError(
-                f"Placement pool returned an invalid layout for env {cur_env} "
-                f"(loss={result.final_loss:.6f}); refusing to write it to simulation."
-            )
-
     anchor_objects_set = set(get_anchor_objects(objects))
     rotations = {obj: get_rotation_xyzw(obj) for obj in objects if obj not in anchor_objects_set}
 
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 27e65803c..bcef88e18 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -129,9 +129,9 @@ def _solve_and_store(self, num_layouts: int) -> None:
     def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
         """Solve layouts that can be used by any env pool.
 
-        Invalid candidates are discarded. If the solver cannot produce enough
-        valid layouts after the bounded refill attempts, ``_solve_and_store``
-        raises instead of letting invalid placements reach simulation.
+        Invalid candidates are discarded when at least one valid layout exists.
+        If no candidate passes strict validation, fall back to best-loss results
+        to match the pre-pool behavior used by existing environments.
         """
         with torch.inference_mode(False):
             result = self._placer.place(self._objects, num_envs=num_layouts, result_per_env=True)
@@ -146,7 +146,11 @@ def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
                 f" {len(valid_layouts)} valid, {num_layouts - len(valid_layouts)} failed validation"
             )
 
-        return valid_layouts
+        if valid_layouts:
+            return valid_layouts
+
+        print("Warning: No candidates passed strict validation. Accepting best-loss layouts as fallback.")
+        return all_layouts
 
     def _store_reusable_results(self, layouts: list[PlacementResult]) -> None:
         """Distribute reusable layouts across env pools using greedy shortest-first.
@@ -199,23 +203,33 @@ def _solve_layouts_with_env_bboxes(self, num_layouts: int) -> tuple[list[Placeme
         return all_results, layouts_per_env
 
     def _store_env_matched_results(self, all_results: list[PlacementResult], layouts_per_env: int) -> None:
-        """Store only successful env-matched results into their corresponding pools."""
+        """Store env-matched results into their corresponding pools.
+
+        Prefer successful layouts for each env. If a specific env has no valid
+        layouts in the batch, fall back to its best-loss results so existing
+        environments with imperfect validation can still run.
+        """
         total_valid = 0
-        for i, r in enumerate(all_results):
-            cur_env = i // layouts_per_env
-            if r.success:
-                self._layout_pools[cur_env].append(r)
-                total_valid += 1
+        fallback_envs = []
+        for cur_env in range(self._num_envs):
+            start = cur_env * layouts_per_env
+            env_results = all_results[start : start + layouts_per_env]
+            valid_results = [r for r in env_results if r.success]
+            if valid_results:
+                self._layout_pools[cur_env].extend(valid_results)
+                total_valid += len(valid_results)
+            else:
+                self._layout_pools[cur_env].extend(env_results)
+                fallback_envs.append(cur_env)
 
         total_solved = len(all_results)
         if total_valid < total_solved:
-            failed_envs = [cur_env for cur_env in self._layout_pools if not self._layout_pools[cur_env]]
             msg = (
                 f"Placement pool (env-specific bbox layouts): solved {total_solved} candidates,"
                 f" {total_valid} valid, {total_solved - total_valid} failed validation"
             )
-            if failed_envs:
-                msg += f". Envs with zero valid layouts: {failed_envs}"
+            if fallback_envs:
+                msg += f". Falling back to best-loss layouts for envs: {fallback_envs}"
             print(msg)
 
     # ------------------------------------------------------------------
diff --git a/isaaclab_arena/tests/test_no_collision_loss.py b/isaaclab_arena/tests/test_no_collision_loss.py
index a52ad09bb..0f2369157 100644
--- a/isaaclab_arena/tests/test_no_collision_loss.py
+++ b/isaaclab_arena/tests/test_no_collision_loss.py
@@ -3,9 +3,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-# pyright: reportPrivateUsage=false
-# This file intentionally probes RelationSolver internals for focused loss checks.
-
 """Tests for NoCollisionLossStrategy and RelationSolver built-in no-overlap behavior."""
 
 import torch
@@ -240,7 +237,7 @@ def test_no_overlap_skips_direct_on_non_anchor_pair():
     state = RelationSolverState(objects, initial_positions, device=torch.device("cpu"))
     solver = RelationSolver(params=RelationSolverParams(max_iters=0))
 
-    loss = solver._compute_no_overlap_loss(state)
+    loss = solver._compute_no_overlap_loss(state)  # pyright: ignore[reportPrivateUsage]
 
     assert torch.isfinite(loss).all()
     assert torch.allclose(loss, torch.zeros_like(loss))
diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py
index 505e4c9ec..f7c8a004b 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -8,8 +8,6 @@
 import torch
 from unittest.mock import MagicMock
 
-import pytest
-
 
 def _create_test_objects():
     """Create a desk (anchor) with two boxes (On + NextTo)."""
@@ -274,8 +272,8 @@ def test_solve_and_place_objects_partial_reset_reusable_pool_consumes_only_reset
     assert available_before - available_after == len(env_ids)
 
 
-def test_solve_and_place_objects_rejects_invalid_pool_layout():
-    """Invalid pool layouts should not be written to simulation."""
+def test_solve_and_place_objects_writes_invalid_fallback_layout():
+    """Invalid fallback layouts should still be written, matching pool fallback behavior."""
 
     from isaaclab_arena.relations.placement_events import solve_and_place_objects
     from isaaclab_arena.relations.placement_result import PlacementResult
@@ -298,46 +296,11 @@ def sample_without_replacement(self, count: int) -> list[PlacementResult]:
                 )
             ]
 
-    with pytest.raises(RuntimeError, match="invalid layout"):
-        solve_and_place_objects(env, torch.tensor([0]), objects, InvalidPool())
-
-    assert len(env._assets) == 0
+    solve_and_place_objects(env, torch.tensor([0]), objects, InvalidPool())
 
-
-def test_solve_and_place_objects_rejects_invalid_layout_before_partial_write():
-    """Invalid layouts should abort the whole reset before any env is written."""
-
-    from isaaclab_arena.relations.placement_events import solve_and_place_objects
-    from isaaclab_arena.relations.placement_result import PlacementResult
-
-    desk, box1, box2 = _create_test_objects()
-    objects = [desk, box1, box2]
-    env = _make_mock_env(num_envs=2)
-
-    class PartiallyInvalidPool:
-        requires_env_indexed_layouts = False
-
-        def sample_without_replacement(self, count: int) -> list[PlacementResult]:
-            assert count == 2
-            return [
-                PlacementResult(
-                    success=True,
-                    positions={box1: (0.1, 0.1, 0.1), box2: (0.2, 0.2, 0.2)},
-                    final_loss=0.0,
-                    attempts=1,
-                ),
-                PlacementResult(
-                    success=False,
-                    positions={box1: (0.0, 0.0, 0.0), box2: (0.0, 0.0, 0.0)},
-                    final_loss=float("nan"),
-                    attempts=1,
-                ),
-            ]
-
-    with pytest.raises(RuntimeError, match="invalid layout"):
-        solve_and_place_objects(env, torch.tensor([0, 1]), objects, PartiallyInvalidPool())
-
-    assert len(env._assets) == 0
+    assert set(env._assets) == {box1.name, box2.name}
+    assert env._assets[box1.name].write_root_pose_to_sim.call_count == 1
+    assert env._assets[box2.name].write_root_pose_to_sim.call_count == 1
 
 
 def test_pooled_placer_sample_without_replacement_returns_different_layouts():
@@ -438,8 +401,8 @@ def test_resolve_on_reset_false_applies_pose_per_env():
             assert p.position_xyz is not None, f"Position should not be None for {obj.name}"
 
 
-def test_pooled_placer_raises_when_no_valid_layouts():
-    """PooledObjectPlacer should fail instead of storing invalid layouts."""
+def test_pooled_placer_falls_back_when_no_valid_layouts():
+    """PooledObjectPlacer should keep best-loss fallback layouts when validation rejects all candidates."""
 
     from isaaclab_arena.assets.dummy_object import DummyObject
     from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
@@ -471,5 +434,7 @@ def test_pooled_placer_raises_when_no_valid_layouts():
     solver_params = RelationSolverParams(max_iters=50, convergence_threshold=1e-6)
     placer_params = ObjectPlacerParams(solver_params=solver_params, max_placement_attempts=1)
 
-    with pytest.raises(RuntimeError, match="could not fill"):
-        PooledObjectPlacer(objects=[desk, big1, big2], placer_params=placer_params, pool_size=5)
+    pool = PooledObjectPlacer(objects=[desk, big1, big2], placer_params=placer_params, pool_size=5)
+
+    assert pool.remaining == 5
+    assert not pool.sample_without_replacement(1)[0].success
diff --git a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
index 466c71b7e..e8bf8ab87 100644
--- a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
+++ b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
@@ -204,6 +204,8 @@ def _build_heterogeneous_objects(self, tabletop_reference, object_names=None):
         from isaaclab_arena.assets.object_set import RigidObjectSet
         from isaaclab_arena.relations.relations import AtPosition, On
 
+        # TODO(@zhx06): Address residual object bouncing with xy-only no-collision
+        # constraints and anchor constraint handling in a follow-up change.
         if object_names:
             print(
                 "Warning: --objects with --mode heterogeneous wraps each object as a "

From 50cf8cdd6162d523f69dbc4eb5224ec05e3933ed Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 13 May 2026 13:24:37 -0700
Subject: [PATCH 30/46] address comments

---
 isaaclab_arena/assets/object_set.py           |  5 ++--
 isaaclab_arena/relations/placement_events.py  |  6 ++++-
 .../relations/pooled_object_placer.py         | 13 ++++++++++
 isaaclab_arena/tests/test_object_set.py       |  5 ++++
 isaaclab_arena/tests/test_placement_events.py | 24 ++++++++++++++++++-
 5 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index dd053f6d1..7bd2acb49 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -151,8 +151,8 @@ def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
 
     def get_contact_sensor_cfg(self, contact_against_object: ObjectBase | None = None) -> ContactSensorCfg:
         # We assume that by here, our USDs have been modified to be compatible with each other
-        # and we can use the first USD path to find the shallowest rigid body.
-        return super().get_contact_sensor_cfg(contact_against_object, usd_path=self.object_usd_paths[0])
+        # and we can use the canonical first member USD to find the shallowest rigid body.
+        return super().get_contact_sensor_cfg(contact_against_object, usd_path=self._member_object_usd_paths[0])
 
     def _generate_variant_indices(self, num_envs: int) -> list[int]:
         n = len(self.objects)
@@ -171,6 +171,7 @@ def _set_variant_indices_by_env(self, variant_indices_by_env: list[int]) -> None
 
         self.variant_indices_by_env = list(variant_indices_by_env)
         if len(self.objects) > 1:
+            # Keep Isaac Lab's MultiUsdFileCfg aligned with the fixed per-env variant assignment.
             self.object_usd_paths = [self._member_object_usd_paths[idx] for idx in self.variant_indices_by_env]
             spawn_cfg = self.object_cfg.spawn if getattr(self, "object_cfg", None) is not None else None
             if isinstance(spawn_cfg, sim_utils.MultiUsdFileCfg):
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index cf0a85da3..4ceb3d139 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -50,7 +50,11 @@ def solve_and_place_objects(
 
     reset_env_ids = env_ids.tolist()
     if placement_pool.requires_env_indexed_layouts:
-        all_results = placement_pool.sample_without_replacement(env.scene.env_origins.shape[0])
+        num_scene_envs = env.scene.env_origins.shape[0]
+        assert (
+            placement_pool.num_envs == num_scene_envs
+        ), f"Placement pool has {placement_pool.num_envs} envs, but scene has {num_scene_envs} env origins."
+        all_results = placement_pool.sample_without_replacement(num_scene_envs)
         results_by_env = {cur_env: all_results[cur_env] for cur_env in reset_env_ids}
     else:
         reset_results = placement_pool.sample_without_replacement(len(reset_env_ids))
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index bcef88e18..f7852df0b 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -63,6 +63,7 @@ def __init__(
         self._objects = list(objects)
         self._placer = ObjectPlacer(params=placer_params)
         self._pool_size = pool_size
+        self._had_fallbacks = False
         self._layout_pools: dict[int, list[PlacementResult]] = {cur_env: [] for cur_env in range(self._num_envs)}
         self._layout_cursors: dict[int, int] = {cur_env: 0 for cur_env in range(self._num_envs)}
 
@@ -149,6 +150,7 @@ def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
         if valid_layouts:
             return valid_layouts
 
+        self._had_fallbacks = True
         print("Warning: No candidates passed strict validation. Accepting best-loss layouts as fallback.")
         return all_layouts
 
@@ -221,6 +223,7 @@ def _store_env_matched_results(self, all_results: list[PlacementResult], layouts
             else:
                 self._layout_pools[cur_env].extend(env_results)
                 fallback_envs.append(cur_env)
+                self._had_fallbacks = True
 
         total_solved = len(all_results)
         if total_valid < total_solved:
@@ -308,6 +311,16 @@ def requires_env_indexed_layouts(self) -> bool:
         """Whether sampled layouts must be matched back to absolute env ids."""
         return self._uses_env_specific_bboxes
 
+    @property
+    def num_envs(self) -> int:
+        """Number of environment pools managed by this placer."""
+        return self._num_envs
+
+    @property
+    def had_fallbacks(self) -> bool:
+        """Whether any pool refill accepted best-loss layouts that failed strict validation."""
+        return self._had_fallbacks
+
     def sample_with_replacement(self, count: int) -> list[PlacementResult]:
         """Pick *count* layouts at random per env-slot (non-consuming).
 
diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py
index 9a8211b6a..b32dc3e27 100644
--- a/isaaclab_arena/tests/test_object_set.py
+++ b/isaaclab_arena/tests/test_object_set.py
@@ -55,6 +55,11 @@ def _test_object_set_samples_and_stores_variant_indices(simulation_app):
     assert getattr(spawn_cfg, "usd_path") == obj_set.object_usd_paths
     assert getattr(spawn_cfg, "random_choice") is False
 
+    with patch("isaaclab_arena.assets.object.find_shallowest_rigid_body", return_value="/rigid") as find_rigid_body:
+        contact_sensor_cfg = obj_set.get_contact_sensor_cfg()
+    find_rigid_body.assert_called_once_with(can_a.usd_path, relative_to_root=True)
+    assert contact_sensor_cfg.prim_path == f"{obj_set.prim_path}/rigid"
+
     per_env_bbox = obj_set.get_bounding_box_per_env(num_envs=4)
     assert torch.allclose(per_env_bbox.max_point[0], bbox_b.max_point[0])
     assert torch.allclose(per_env_bbox.max_point[1], bbox_a.max_point[0])
diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py
index f7c8a004b..947679191 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -8,6 +8,8 @@
 import torch
 from unittest.mock import MagicMock
 
+import pytest
+
 
 def _create_test_objects():
     """Create a desk (anchor) with two boxes (On + NextTo)."""
@@ -303,6 +305,23 @@ def sample_without_replacement(self, count: int) -> list[PlacementResult]:
     assert env._assets[box2.name].write_root_pose_to_sim.call_count == 1
 
 
+def test_solve_and_place_objects_asserts_env_indexed_pool_size_matches_scene():
+    """Env-indexed pool slots must line up with absolute Isaac Lab env ids."""
+
+    from isaaclab_arena.relations.placement_events import solve_and_place_objects
+
+    desk, box1, box2 = _create_test_objects()
+    objects = [desk, box1, box2]
+    env = _make_mock_env(num_envs=2)
+
+    class MismatchedEnvIndexedPool:
+        requires_env_indexed_layouts = True
+        num_envs = 1
+
+    with pytest.raises(AssertionError, match="scene has 2 env origins"):
+        solve_and_place_objects(env, torch.tensor([0]), objects, MismatchedEnvIndexedPool())
+
+
 def test_pooled_placer_sample_without_replacement_returns_different_layouts():
     """sample_without_replacement() should return layouts (likely different across draws)."""
 
@@ -401,7 +420,7 @@ def test_resolve_on_reset_false_applies_pose_per_env():
             assert p.position_xyz is not None, f"Position should not be None for {obj.name}"
 
 
-def test_pooled_placer_falls_back_when_no_valid_layouts():
+def test_pooled_placer_falls_back_when_no_valid_layouts(capsys):
     """PooledObjectPlacer should keep best-loss fallback layouts when validation rejects all candidates."""
 
     from isaaclab_arena.assets.dummy_object import DummyObject
@@ -435,6 +454,9 @@ def test_pooled_placer_falls_back_when_no_valid_layouts():
     placer_params = ObjectPlacerParams(solver_params=solver_params, max_placement_attempts=1)
 
     pool = PooledObjectPlacer(objects=[desk, big1, big2], placer_params=placer_params, pool_size=5)
+    captured = capsys.readouterr()
 
     assert pool.remaining == 5
+    assert pool.had_fallbacks
+    assert "Accepting best-loss layouts as fallback" in captured.out
     assert not pool.sample_without_replacement(1)[0].success

From e2ec342a97345b8fad5ef510edc24257c0e17978 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Fri, 15 May 2026 09:48:41 -0700
Subject: [PATCH 31/46] add bbox helper, edit comments

---
 isaaclab_arena/assets/dummy_object.py         |  9 ---
 isaaclab_arena/assets/object_base.py          | 23 -------
 .../isaaclab_arena_environment.py             |  4 +-
 isaaclab_arena/relations/bbox_helpers.py      | 64 +++++++++++++++++++
 isaaclab_arena/relations/object_placer.py     |  9 +--
 .../relations/pooled_object_placer.py         | 15 ++---
 isaaclab_arena/relations/relation_solver.py   | 40 ++++--------
 .../relations/relation_solver_state.py        | 17 +++++
 .../tests/test_heterogeneous_placement.py     |  9 +--
 9 files changed, 111 insertions(+), 79 deletions(-)
 create mode 100644 isaaclab_arena/relations/bbox_helpers.py

diff --git a/isaaclab_arena/assets/dummy_object.py b/isaaclab_arena/assets/dummy_object.py
index 009788f47..cb7ab52a3 100644
--- a/isaaclab_arena/assets/dummy_object.py
+++ b/isaaclab_arena/assets/dummy_object.py
@@ -25,7 +25,6 @@ def __init__(
         self.name = name
         self.initial_pose = initial_pose
         self.bounding_box = bounding_box
-        self.has_env_specific_bboxes = False
         assert self.bounding_box is not None
         self.relations = list(relations or [])
 
@@ -43,14 +42,6 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get local bounding box (relative to object origin)."""
         return self.bounding_box
 
-    def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
-        """Mirror ObjectBase.get_bounding_box_per_env for this test double."""
-        bbox = self.get_bounding_box()
-        return AxisAlignedBoundingBox(
-            min_point=bbox.min_point.expand(num_envs, 3),
-            max_point=bbox.max_point.expand(num_envs, 3),
-        )
-
     def get_world_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get bounding box in world coordinates (local bbox rotated and translated).
 
diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py
index cec3f6ecb..075d8f1c9 100644
--- a/isaaclab_arena/assets/object_base.py
+++ b/isaaclab_arena/assets/object_base.py
@@ -37,8 +37,6 @@
 class ObjectBase(Asset, ABC):
     """Parent class for (spawnable) Object and ObjectReference."""
 
-    has_env_specific_bboxes: bool = False
-
     def __init__(
         self,
         name: str,
@@ -75,27 +73,6 @@ def get_world_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get bounding box in world coordinates (local bbox rotated and translated)."""
         ...
 
-    def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
-        """Get local bounding boxes for each environment.
-
-        This default implementation is for objects with the same geometry in
-        every environment: it expands the single local bbox to ``(num_envs, 3)``.
-        ``RigidObjectSet`` overrides this to return the bbox for each env's
-        assigned variant.
-
-        Args:
-            num_envs: Number of environments.
-
-        Returns:
-            ``AxisAlignedBoundingBox`` with ``min_point`` / ``max_point`` of shape
-            ``(num_envs, 3)``.
-        """
-        bbox = self.get_bounding_box()
-        return AxisAlignedBoundingBox(
-            min_point=bbox.min_point.expand(num_envs, 3),
-            max_point=bbox.max_point.expand(num_envs, 3),
-        )
-
     def _get_initial_pose_as_pose(self) -> Pose | None:
         """Return a single ``Pose`` suitable for *init_state* and bounding-box calculations.
 
diff --git a/isaaclab_arena/environments/isaaclab_arena_environment.py b/isaaclab_arena/environments/isaaclab_arena_environment.py
index b481cab0a..c8f2a9769 100644
--- a/isaaclab_arena/environments/isaaclab_arena_environment.py
+++ b/isaaclab_arena/environments/isaaclab_arena_environment.py
@@ -26,9 +26,7 @@ def __init__(
         embodiment: EmbodimentBase | None = None,
         task: TaskBase | None = None,
         teleop_device: TeleopDeviceBase | None = None,
-        env_cfg_callback: (
-            Callable[[IsaacLabArenaManagerBasedRLEnvCfg], IsaacLabArenaManagerBasedRLEnvCfg] | None
-        ) = None,
+        env_cfg_callback: Callable[IsaacLabArenaManagerBasedRLEnvCfg] | None = None,
         rl_framework_entry_point: str | None = None,
         rl_policy_cfg: str | None = None,
     ):
diff --git a/isaaclab_arena/relations/bbox_helpers.py b/isaaclab_arena/relations/bbox_helpers.py
new file mode 100644
index 000000000..f1fa04047
--- /dev/null
+++ b/isaaclab_arena/relations/bbox_helpers.py
@@ -0,0 +1,64 @@
+# 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
+
+"""Bounding-box utilities for the relation solver / object placer.
+
+These helpers centralise the logic that expands single-object bounding boxes
+to per-environment tensors and detects whether any object carries
+env-specific geometry (e.g. ``RigidObjectSet`` with multiple variants).
+
+Placing this logic here (instead of on ``ObjectBase``) keeps ``num_envs``
+and other placement/simulation details out of the asset layer.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+
+if TYPE_CHECKING:
+    from isaaclab_arena.assets.object_base import ObjectBase
+
+
+def object_has_env_specific_bboxes(obj: ObjectBase) -> bool:
+    """Whether *obj* produces different bounding boxes per environment.
+
+    Returns ``True`` for multi-variant ``RigidObjectSet`` instances (which
+    override ``get_bounding_box_per_env``). All other objects return ``False``.
+    """
+    return getattr(obj, "has_env_specific_bboxes", False)
+
+
+def any_object_has_env_specific_bboxes(objects: list[ObjectBase]) -> bool:
+    """Whether any object in *objects* has env-specific bounding boxes."""
+    return any(object_has_env_specific_bboxes(obj) for obj in objects)
+
+
+def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBoundingBox:
+    """Return per-environment local bounding boxes for *obj*.
+
+    * For objects with env-specific variants (``RigidObjectSet``), delegates
+      to the object's own ``get_bounding_box_per_env`` override.
+    * For all other objects, expands the single local bbox to
+      ``(num_envs, 3)``.
+
+    Args:
+        obj: The object to query.
+        num_envs: Number of environments.
+
+    Returns:
+        ``AxisAlignedBoundingBox`` with ``min_point`` / ``max_point`` of shape
+        ``(num_envs, 3)``.
+    """
+    override = getattr(obj, "get_bounding_box_per_env", None)
+    if override is not None:
+        return override(num_envs)
+
+    bbox = obj.get_bounding_box()
+    return AxisAlignedBoundingBox(
+        min_point=bbox.min_point.expand(num_envs, 3),
+        max_point=bbox.max_point.expand(num_envs, 3),
+    )
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index a7633a636..ddfd72386 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -10,6 +10,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+from isaaclab_arena.relations.bbox_helpers import any_object_has_env_specific_bboxes, get_bounding_box_per_env
 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
@@ -103,7 +104,7 @@ def place(
                 identically to all environments.
             env_bboxes: Pre-computed per-env bounding boxes (shape ``(num_envs, 3)``
                 per object).  When provided, ``_place_heterogeneous`` uses these
-                instead of calling ``get_bounding_box_per_env(num_envs)``.  This
+                instead of calling ``get_bounding_box_per_env(obj, num_envs)``.  This
                 allows the caller to tile real-env bboxes for pooled solving.
 
         Returns:
@@ -148,7 +149,7 @@ def place(
         #   - homogeneous: any solved layout serves any env; pick the top num_results.
         #   - heterogeneous: some objects vary per env, so each env owns a fixed slice
         #     of candidates and we pick the best within that slice.
-        uses_env_specific_bboxes = result_per_env and any(obj.has_env_specific_bboxes for obj in objects)
+        uses_env_specific_bboxes = result_per_env and any_object_has_env_specific_bboxes(objects)
 
         if uses_env_specific_bboxes:
             results_per_env = self._place_heterogeneous(
@@ -251,10 +252,10 @@ def _place_heterogeneous(
         Args:
             env_bboxes: Optional pre-tiled per-env bboxes of shape
                 ``(num_envs, 3)``. When ``None`` we call
-                ``get_bounding_box_per_env(num_envs)`` on each object.
+                ``get_bounding_box_per_env(obj, num_envs)`` on each object.
         """
         if env_bboxes is None:
-            env_bboxes = {obj: obj.get_bounding_box_per_env(num_envs) for obj in objects}
+            env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
 
         # Build the per-env (1, 3) bbox views once: used both for On-guided
         # initial-position sampling and for per-env validation below.
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index f7852df0b..8f5a959cb 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -9,6 +9,7 @@
 import torch
 from typing import TYPE_CHECKING
 
+from isaaclab_arena.relations.bbox_helpers import any_object_has_env_specific_bboxes, get_bounding_box_per_env
 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, PlacementResult
@@ -22,8 +23,7 @@ class PooledObjectPlacer:
     """Object placer that keeps a pool of optimized layouts.
 
     Storage: ``num_envs`` independent layout pools, each with its own read
-    cursor (this replaces the single ``_layouts`` list + ``_next_idx`` cursor
-    used before heterogeneous placement). Env-specific layouts are solved
+    cursor. Env-specific layouts are solved
     against a fixed env's object geometry and must be sampled in complete env
     rounds. Reusable layouts can be consumed one at a time.
 
@@ -52,7 +52,7 @@ def __init__(
         # 1. Validate params.
         if pool_size < 1:
             raise ValueError(f"pool_size must be >= 1, got {pool_size}")
-        self._uses_env_specific_bboxes = any(obj.has_env_specific_bboxes for obj in objects)
+        self._uses_env_specific_bboxes = any_object_has_env_specific_bboxes(objects)
         if self._uses_env_specific_bboxes:
             assert num_envs is not None, "num_envs is required when layouts use env-specific object variants."
         self._num_envs = num_envs if num_envs is not None else 1
@@ -182,7 +182,7 @@ def _solve_layouts_with_env_bboxes(self, num_layouts: int) -> tuple[list[Placeme
         layouts_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
         total_layouts = layouts_per_env * self._num_envs
 
-        real_bboxes = {obj: obj.get_bounding_box_per_env(self._num_envs) for obj in self._objects}
+        real_bboxes = {obj: get_bounding_box_per_env(obj, self._num_envs) for obj in self._objects}
 
         # (num_envs, 3) -> repeat each env's row layouts_per_env times -> (total_layouts, 3).
         tiled_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {
@@ -341,10 +341,9 @@ def sample_with_replacement(self, count: int) -> list[PlacementResult]:
     def remaining(self) -> int:
         """Number of complete env rounds available to :meth:`sample_without_replacement`.
 
-        Returns the minimum unread count across env pools (the previous
-        ``remaining`` was a total across one shared list; under per-env
-        storage a single round consumes one layout from every env, so the
-        minimum is what limits without-replacement capacity).
+        Returns the minimum unread count across env pools. A single round
+        consumes one layout from every env, so the minimum is what limits
+        without-replacement capacity.
         """
         return min(self._available_per_env())
 
diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index a70055356..2a23b0f1f 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -67,32 +67,17 @@ def _get_strategy(self, relation: RelationBase) -> RelationLossStrategy | UnaryR
             )
         return strategy
 
-    def _get_bbox(
-        self,
-        obj: ObjectBase,
-        device: torch.device | None,
-        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None,
-    ) -> AxisAlignedBoundingBox:
-        """Return the per-env or default local bbox for *obj*, moved to *device*."""
-        if env_bboxes is not None and obj in env_bboxes:
-            return env_bboxes[obj].to(device)
-        return obj.get_bounding_box().to(device)
-
     def _compute_total_loss(
         self,
         state: RelationSolverState,
         debug: bool = False,
-        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> torch.Tensor:
         """Compute total loss from all relations using registered strategies.
 
         Args:
-            state: Current optimization state with object positions.
+            state: Current optimization state with object positions and
+                optional per-env bounding boxes (accessed via ``state.get_bbox``).
             debug: If True, print detailed loss breakdown.
-            env_bboxes: Optional per-env bboxes keyed by object. When
-                provided the bbox ``min_point`` / ``max_point`` have shape
-                ``(batch, 3)`` instead of ``(1, 3)``, enabling heterogeneous
-                object placement.
 
         Returns:
             Scalar loss tensor (mean over environments).
@@ -106,7 +91,7 @@ def _compute_total_loss(
             for relation in obj.get_spatial_relations():
                 child_pos = state.get_position(obj)
                 strategy = self._get_strategy(relation)
-                child_bbox = self._get_bbox(obj, device, env_bboxes)
+                child_bbox = state.get_bbox(obj)
 
                 # Handle unary relations (no parent)
                 if isinstance(relation, UnaryRelation):
@@ -126,7 +111,7 @@ def _compute_total_loss(
                         parent_world_bbox = parent.get_world_bounding_box().to(device)
                     else:
                         parent_pos = state.get_position(parent)
-                        parent_bbox = self._get_bbox(parent, device, env_bboxes)
+                        parent_bbox = state.get_bbox(parent)
                         parent_world_bbox = parent_bbox.translated(parent_pos)
                     loss = relation_strategy.compute_loss(
                         relation=relation,
@@ -143,7 +128,7 @@ def _compute_total_loss(
                 total_loss = total_loss + loss
 
         # Add built-in no-overlap loss between all object pairs
-        total_loss = total_loss + self._compute_no_overlap_loss(state, debug, env_bboxes=env_bboxes)
+        total_loss = total_loss + self._compute_no_overlap_loss(state, debug)
 
         self._last_loss_per_env = total_loss.detach().clone()
         return total_loss.mean()
@@ -152,7 +137,6 @@ def _compute_no_overlap_loss(
         self,
         state: RelationSolverState,
         debug: bool = False,
-        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> torch.Tensor:
         """Compute pairwise no-overlap loss for all non-anchor objects against all other objects.
 
@@ -162,9 +146,9 @@ def _compute_no_overlap_loss(
           the loss in both directions with the other's position detached.
 
         Args:
-            state: Current optimization state with object positions.
+            state: Current optimization state with object positions and
+                optional per-env bounding boxes.
             debug: If True, print detailed loss breakdown.
-            env_bboxes: Optional per-env bboxes for heterogeneous placement.
 
         Returns:
             Per-environment loss tensor of shape (batch_size,).
@@ -183,7 +167,7 @@ def _compute_no_overlap_loss(
 
         for i, child in enumerate(non_anchor_objects):
             child_pos = state.get_position(child)
-            child_bbox = self._get_bbox(child, device, env_bboxes)
+            child_bbox = state.get_bbox(child)
 
             # Against all anchors
             for anchor in anchor_objects:
@@ -206,7 +190,7 @@ def _compute_no_overlap_loss(
                 if (id(child), id(other)) in on_pairs:
                     continue
                 other_pos = state.get_position(other)
-                other_bbox = self._get_bbox(other, device, env_bboxes)
+                other_bbox = state.get_bbox(other)
 
                 # Forward: gradient flows to child (object i)
                 other_world_bbox = other_bbox.translated(other_pos.detach())
@@ -258,7 +242,7 @@ def solve(
             List of dicts (one per env) mapping objects to their solved (x, y, z) positions.
         """
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        state = RelationSolverState(objects, initial_positions, device=device)
+        state = RelationSolverState(objects, initial_positions, device=device, env_bboxes=env_bboxes)
 
         if self.params.verbose:
             anchor_names = [obj.name for obj in state.anchor_objects]
@@ -282,7 +266,7 @@ def solve(
         # Compute initial loss so _last_loss_per_env is always populated
         # (needed even when max_iters=0, e.g. tests that only check init positions).
         with torch.no_grad():
-            self._compute_total_loss(state, env_bboxes=env_bboxes)
+            self._compute_total_loss(state)
 
         # Optimization loop
         loss_history = []
@@ -295,7 +279,7 @@ def solve(
                 position_history.append(state.get_all_positions_snapshot())
 
             # Compute total loss
-            loss = self._compute_total_loss(state, env_bboxes=env_bboxes)
+            loss = self._compute_total_loss(state)
             loss_history.append(loss.item())
 
             # Backprop and update (only optimizable positions will update)
diff --git a/isaaclab_arena/relations/relation_solver_state.py b/isaaclab_arena/relations/relation_solver_state.py
index 344185fed..3d837a377 100644
--- a/isaaclab_arena/relations/relation_solver_state.py
+++ b/isaaclab_arena/relations/relation_solver_state.py
@@ -9,6 +9,7 @@
 from typing import TYPE_CHECKING
 
 from isaaclab_arena.relations.relations import get_anchor_objects
+from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
     from isaaclab_arena.assets.object_base import ObjectBase
@@ -29,6 +30,7 @@ def __init__(
         objects: list[ObjectBase],
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]],
         device: torch.device | None = None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ):
         """Initialize optimization state.
 
@@ -38,6 +40,9 @@ def __init__(
             initial_positions: List of dicts (one per env). Length 1 = single-env,
                 length > 1 = batched.
             device: Torch device for all tensors. Defaults to CPU.
+            env_bboxes: Optional per-env bounding boxes keyed by object.
+                When provided, ``get_bbox(obj)`` returns the per-env bbox
+                (shape ``(batch, 3)``) instead of the object's default.
         """
         assert len(initial_positions) >= 1, "initial_positions must contain at least one dict."
         anchor_objects = get_anchor_objects(objects)
@@ -93,6 +98,8 @@ def __init__(
             self._opt_idx_tensor = None
             self._optimizable_positions = None
 
+        self._env_bboxes = env_bboxes
+
     @property
     def device(self) -> torch.device:
         """Torch device for all position tensors."""
@@ -142,6 +149,16 @@ def get_position(self, obj: ObjectBase) -> torch.Tensor:
         opt_idx = self._global_to_opt_idx[idx]
         return self._optimizable_positions[:, opt_idx, :]
 
+    def get_bbox(self, obj: ObjectBase) -> AxisAlignedBoundingBox:
+        """Return the local bounding box for *obj*, moved to the state's device.
+
+        Uses the per-env override from ``env_bboxes`` when available, otherwise
+        falls back to ``obj.get_bounding_box()``.
+        """
+        if self._env_bboxes is not None and obj in self._env_bboxes:
+            return self._env_bboxes[obj].to(self._device)
+        return obj.get_bounding_box().to(self._device)
+
     def get_all_positions_snapshot(self) -> list[tuple[float, float, float]]:
         """Get detached copy of all positions for history tracking.
 
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 82e822bb7..07abb1173 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -14,6 +14,7 @@
 import pytest
 
 from isaaclab_arena.assets.dummy_object import DummyObject
+from isaaclab_arena.relations.bbox_helpers import get_bounding_box_per_env, object_has_env_specific_bboxes
 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, PlacementResult
@@ -73,7 +74,7 @@ def test_dummy_object_bbox_per_env_expands_single():
         bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)),
     )
 
-    per_env = obj.get_bounding_box_per_env(4)
+    per_env = get_bounding_box_per_env(obj, 4)
     assert per_env.min_point.shape == (4, 3)
     assert per_env.max_point.shape == (4, 3)
     assert torch.allclose(per_env.min_point[0], per_env.min_point[3])
@@ -86,7 +87,7 @@ def test_heterogeneous_dummy_returns_different_bboxes():
     large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3))
     obj = HeterogeneousDummyObject(name="set", bboxes=[small, large])
 
-    per_env = obj.get_bounding_box_per_env(4)
+    per_env = get_bounding_box_per_env(obj, 4)
     assert per_env.max_point.shape == (4, 3)
     # env 0 and 2 should use small; env 1 and 3 should use large
     assert torch.allclose(per_env.max_point[0], torch.tensor([0.1, 0.1, 0.1]))
@@ -104,7 +105,7 @@ def test_dummy_object_preserves_constructor_relations():
     )
 
     assert obj.get_relations() == [anchor_relation]
-    assert obj.has_env_specific_bboxes is False
+    assert object_has_env_specific_bboxes(obj) is False
 
 
 def test_object_preserves_constructor_relations():
@@ -122,7 +123,7 @@ def test_object_preserves_constructor_relations():
     )
 
     assert obj.get_relations() == [anchor_relation]
-    assert obj.has_env_specific_bboxes is False
+    assert object_has_env_specific_bboxes(obj) is False
 
 
 # ---------------------------------------------------------------------------

From 50c94cf14eedfc7f2a52c4d771e5a240c8b5850c Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Fri, 15 May 2026 10:58:00 -0700
Subject: [PATCH 32/46] edit comments style

---
 isaaclab_arena/relations/bbox_helpers.py      | 24 +++++-----
 isaaclab_arena/relations/object_placer.py     | 36 +++++++--------
 isaaclab_arena/relations/placement_events.py  |  2 +-
 .../relations/pooled_object_placer.py         | 44 +++++++++----------
 isaaclab_arena/relations/relation_solver.py   | 10 ++---
 .../relations/relation_solver_state.py        | 10 ++---
 6 files changed, 62 insertions(+), 64 deletions(-)

diff --git a/isaaclab_arena/relations/bbox_helpers.py b/isaaclab_arena/relations/bbox_helpers.py
index f1fa04047..b4a890aa2 100644
--- a/isaaclab_arena/relations/bbox_helpers.py
+++ b/isaaclab_arena/relations/bbox_helpers.py
@@ -7,9 +7,9 @@
 
 These helpers centralise the logic that expands single-object bounding boxes
 to per-environment tensors and detects whether any object carries
-env-specific geometry (e.g. ``RigidObjectSet`` with multiple variants).
+env-specific geometry (e.g. RigidObjectSet with multiple variants).
 
-Placing this logic here (instead of on ``ObjectBase``) keeps ``num_envs``
+Placing this logic here (instead of on ObjectBase) keeps num_envs
 and other placement/simulation details out of the asset layer.
 """
 
@@ -24,34 +24,32 @@
 
 
 def object_has_env_specific_bboxes(obj: ObjectBase) -> bool:
-    """Whether *obj* produces different bounding boxes per environment.
+    """Whether obj produces different bounding boxes per environment.
 
-    Returns ``True`` for multi-variant ``RigidObjectSet`` instances (which
-    override ``get_bounding_box_per_env``). All other objects return ``False``.
+    Returns True for multi-variant RigidObjectSet instances (which
+    override get_bounding_box_per_env). All other objects return False.
     """
     return getattr(obj, "has_env_specific_bboxes", False)
 
 
 def any_object_has_env_specific_bboxes(objects: list[ObjectBase]) -> bool:
-    """Whether any object in *objects* has env-specific bounding boxes."""
+    """Whether any object in the list has env-specific bounding boxes."""
     return any(object_has_env_specific_bboxes(obj) for obj in objects)
 
 
 def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBoundingBox:
-    """Return per-environment local bounding boxes for *obj*.
+    """Return per-environment local bounding boxes for obj.
 
-    * For objects with env-specific variants (``RigidObjectSet``), delegates
-      to the object's own ``get_bounding_box_per_env`` override.
-    * For all other objects, expands the single local bbox to
-      ``(num_envs, 3)``.
+    For objects with env-specific variants (RigidObjectSet), delegates to the
+    object's own get_bounding_box_per_env override. For all other objects,
+    expands the single local bbox to (num_envs, 3).
 
     Args:
         obj: The object to query.
         num_envs: Number of environments.
 
     Returns:
-        ``AxisAlignedBoundingBox`` with ``min_point`` / ``max_point`` of shape
-        ``(num_envs, 3)``.
+        AxisAlignedBoundingBox with min_point / max_point of shape (num_envs, 3).
     """
     override = getattr(obj, "get_bounding_box_per_env", None)
     if override is not None:
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index ddfd72386..b90018f98 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -75,10 +75,10 @@ def _resolve_bbox(
         obj: ObjectBase,
         overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None,
     ) -> AxisAlignedBoundingBox:
-        """Return *overrides[obj]* if present, otherwise *obj*'s default bbox.
+        """Return overrides[obj] if present, otherwise the object's default bbox.
 
         Heterogeneous placement passes a single-env override bbox; the
-        homogeneous path and unit tests pass ``None`` and rely on the
+        homogeneous path and unit tests pass None and rely on the
         object's own bbox.
         """
         if overrides is not None and obj in overrides:
@@ -102,9 +102,9 @@ def place(
             result_per_env: When True (default), each environment gets a distinct
                 layout. When False, a single best layout is solved and applied
                 identically to all environments.
-            env_bboxes: Pre-computed per-env bounding boxes (shape ``(num_envs, 3)``
-                per object).  When provided, ``_place_heterogeneous`` uses these
-                instead of calling ``get_bounding_box_per_env(obj, num_envs)``.  This
+            env_bboxes: Pre-computed per-env bounding boxes (shape (num_envs, 3)
+                per object). When provided, _place_heterogeneous uses these
+                instead of calling get_bounding_box_per_env(obj, num_envs). This
                 allows the caller to tile real-env bboxes for pooled solving.
 
         Returns:
@@ -189,8 +189,8 @@ def _place_homogeneous(
     ) -> list[PlacementResult]:
         """Original batched placement path.
 
-        Solves ``max_attempts * num_results`` candidates in one batched
-        ``solver.solve()`` call, then returns the best ``num_results`` ranked
+        Solves max_attempts * num_results candidates in one batched
+        solver.solve() call, then returns the best num_results ranked
         by (is_valid, loss). All envs share the same geometry, so any solved
         layout can serve any env.
         """
@@ -244,15 +244,15 @@ def _place_heterogeneous(
     ) -> list[PlacementResult]:
         """Per-env placement: each candidate is tied to its env's object variants.
 
-        Candidates ``[cur_env * max_attempts : (cur_env + 1) * max_attempts]``
-        belong to ``cur_env`` and use that env's variant geometry. We solve all
-        ``num_envs * max_attempts`` candidates in one batched ``solver.solve()``
+        Candidates [cur_env * max_attempts : (cur_env + 1) * max_attempts]
+        belong to cur_env and use that env's variant geometry. We solve all
+        num_envs * max_attempts candidates in one batched solver.solve()
         call and pick the best candidate within each env's slice.
 
         Args:
             env_bboxes: Optional pre-tiled per-env bboxes of shape
-                ``(num_envs, 3)``. When ``None`` we call
-                ``get_bounding_box_per_env(obj, num_envs)`` on each object.
+                (num_envs, 3). When None we call
+                get_bounding_box_per_env(obj, num_envs) on each object.
         """
         if env_bboxes is None:
             env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
@@ -335,7 +335,7 @@ def _generate_initial_positions(
         Args:
             generator: Optional RNG generator for reproducible sampling. When None,
                 uses PyTorch's global RNG.
-            child_bboxes: Optional per-object bbox overrides with shape ``(1, 3)``.
+            child_bboxes: Optional per-object bbox overrides with shape (1, 3).
                 Used by heterogeneous placement to supply the correct variant
                 bbox when computing On-guided initial positions.
 
@@ -398,8 +398,8 @@ def _compute_on_guided_position(
         Args:
             generator: Optional RNG generator for reproducible sampling. When None,
                 uses PyTorch's global RNG.
-            child_bbox: Optional bbox override for the child object. When ``None``,
-                ``obj.get_bounding_box()`` is used.
+            child_bbox: Optional bbox override for the child object. When None,
+                obj.get_bounding_box() is used.
         """
         on_relation = next(r for r in obj.get_relations() if isinstance(r, On))
         parent_bbox = self._get_on_parent_world_bbox(on_relation.parent, anchor_objects, anchor_bbox)
@@ -469,7 +469,7 @@ def _validate_on_relations(
 
         Args:
             positions: Solved positions for each object.
-            bbox_overrides: Optional per-object bbox overrides (single-env, shape ``(1, 3)``).
+            bbox_overrides: Optional per-object bbox overrides (single-env, shape (1, 3)).
                 Used by heterogeneous placement to supply the correct variant bbox.
         """
         for obj in positions:
@@ -520,7 +520,7 @@ def _validate_no_overlap(
 
         Args:
             positions: Solved positions for each object.
-            bbox_overrides: Optional per-object bbox overrides (single-env, shape ``(1, 3)``).
+            bbox_overrides: Optional per-object bbox overrides (single-env, shape (1, 3)).
                 Used by heterogeneous placement to supply the correct variant bbox.
         """
         # Build set of On-related pairs to skip (child, parent) and (parent, child).
@@ -568,7 +568,7 @@ def _validate_placement(
 
         Args:
             positions: Dictionary mapping objects to their solved (x, y, z) positions.
-            bbox_overrides: Optional per-object bbox overrides (single-env, shape ``(1, 3)``).
+            bbox_overrides: Optional per-object bbox overrides (single-env, shape (1, 3)).
                 Used by heterogeneous placement to supply the correct variant bbox.
 
         Returns:
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index 4ceb3d139..5d94ad72a 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -34,7 +34,7 @@ def solve_and_place_objects(
 ) -> None:
     """Coordinated reset event that draws layouts from the pool and writes poses.
 
-    Registered as a single ``EventTermCfg(mode="reset")``. Env-specific
+    Registered as a single EventTermCfg(mode="reset"). Env-specific
     layouts advance by one full env round so each result still matches its
     absolute env id. Reusable layouts draw only for the environments being
     reset.
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 8f5a959cb..1595333df 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -22,21 +22,21 @@
 class PooledObjectPlacer:
     """Object placer that keeps a pool of optimized layouts.
 
-    Storage: ``num_envs`` independent layout pools, each with its own read
+    Storage: num_envs independent layout pools, each with its own read
     cursor. Env-specific layouts are solved
     against a fixed env's object geometry and must be sampled in complete env
     rounds. Reusable layouts can be consumed one at a time.
 
     The pool is refilled automatically when an env's queue runs out.
 
-    * :meth:`sample_without_replacement` — returns the next *count* layouts.
-      Env-specific layouts require ``count`` to be a multiple of ``num_envs``.
-    * :meth:`sample_with_replacement` — picks *count* layouts at random per
+    * sample_without_replacement — returns the next count layouts.
+      Env-specific layouts require count to be a multiple of num_envs.
+    * sample_with_replacement — picks count layouts at random per
       env-slot (non-consuming). Used for static initial positions.
 
     Args:
         objects: All objects (including anchors) participating in relation solving.
-        placer_params: Parameters forwarded to ``ObjectPlacer`` for the batched solve.
+        placer_params: Parameters forwarded to ObjectPlacer for the batched solve.
         pool_size: Number of layouts to solve per batch.
         num_envs: Total number of simulation environments.  Required when
             layouts use env-specific object variants and defaults to 1 otherwise.
@@ -81,7 +81,7 @@ def __init__(
     # ------------------------------------------------------------------
 
     def _available_per_env(self) -> list[int]:
-        """Number of unread layouts in each env's pool (length ``num_envs``)."""
+        """Number of unread layouts in each env's pool (length num_envs)."""
         return [len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] for cur_env in range(self._num_envs)]
 
     def _total_available(self) -> int:
@@ -96,10 +96,10 @@ def _discard_consumed_layouts(self) -> None:
             self._layout_cursors[cur_env] = 0
 
     def _solve_and_store(self, num_layouts: int) -> None:
-        """Solve layouts in batches until every env has ``target_per_env`` unread layouts.
+        """Solve layouts in batches until every env has target_per_env unread layouts.
 
         Each batch contributes (roughly) one round of layouts per env. The
-        outer loop is bounded by ``max_placement_attempts`` to avoid an
+        outer loop is bounded by max_placement_attempts to avoid an
         unbounded refill in pathological configurations.
         """
         self._discard_consumed_layouts()
@@ -157,7 +157,7 @@ def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
     def _store_reusable_results(self, layouts: list[PlacementResult]) -> None:
         """Distribute reusable layouts across env pools using greedy shortest-first.
 
-        Layouts produced by ``_solve_reusable_layouts`` are interchangeable
+        Layouts produced by _solve_reusable_layouts are interchangeable
         across envs, so we place each one into whichever pool currently has
         the fewest unread layouts. This keeps reusable capacity balanced
         across env pools.
@@ -174,10 +174,10 @@ def _store_reusable_results(self, layouts: list[PlacementResult]) -> None:
     def _solve_layouts_with_env_bboxes(self, num_layouts: int) -> tuple[list[PlacementResult], int]:
         """Solve layouts tied to each env's actual object geometry.
 
-        Computes bounding boxes for the real ``num_envs`` once, tiles them
-        to ``num_layouts`` entries, and solves everything in **one** batched
-        ``place()`` call.  Result ``i`` is mapped back to real env
-        ``i // layouts_per_env`` for pool storage.
+        Computes bounding boxes for the real num_envs once, tiles them
+        to num_layouts entries, and solves everything in one batched
+        place() call. Result i is mapped back to real env
+        i // layouts_per_env for pool storage.
         """
         layouts_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
         total_layouts = layouts_per_env * self._num_envs
@@ -240,19 +240,19 @@ def _store_env_matched_results(self, all_results: list[PlacementResult], layouts
     # ------------------------------------------------------------------
 
     def sample_without_replacement(self, count: int) -> list[PlacementResult]:
-        """Return the next *count* layouts.
+        """Return the next count layouts.
 
         Env-specific layouts are returned as complete rounds of
-        ``[env_0, env_1, ..., env_{num_envs-1}]`` so each result still maps
+        [env_0, env_1, ..., env_{num_envs-1}] so each result still maps
         to the absolute environment it was solved for. Reusable layouts are
-        interchangeable and consume only ``count`` entries.
+        interchangeable and consume only count entries.
 
         Args:
             count: Number of layouts to return.
 
         Raises:
             ValueError: If env-specific layouts are requested without a complete env round.
-            RuntimeError: If the pool cannot provide *count* layouts after refilling.
+            RuntimeError: If the pool cannot provide count layouts after refilling.
         """
         if self._uses_env_specific_bboxes:
             return self._sample_env_indexed_without_replacement(count)
@@ -281,7 +281,7 @@ def _sample_env_indexed_without_replacement(self, count: int) -> list[PlacementR
         return results
 
     def _sample_reusable_without_replacement(self, count: int) -> list[PlacementResult]:
-        """Consume exactly ``count`` interchangeable layouts."""
+        """Consume exactly count interchangeable layouts."""
         if self._total_available() < count:
             self._solve_and_store(max(self._pool_size, count))
 
@@ -322,11 +322,11 @@ def had_fallbacks(self) -> bool:
         return self._had_fallbacks
 
     def sample_with_replacement(self, count: int) -> list[PlacementResult]:
-        """Pick *count* layouts at random per env-slot (non-consuming).
+        """Pick count layouts at random per env-slot (non-consuming).
 
-        Slot ``i`` is filled by a random pick from env ``i % num_envs``'s
-        pool, so a length-``count`` request walks env slots in order. Used
-        by ``resolve_on_reset=False`` to assign initial positions that persist
+        Slot i is filled by a random pick from env i % num_envs's
+        pool, so a length-count request walks env slots in order. Used
+        by resolve_on_reset=False to assign initial positions that persist
         across resets.
         """
         results: list[PlacementResult] = []
diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index 2a23b0f1f..72d037ce2 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -76,7 +76,7 @@ def _compute_total_loss(
 
         Args:
             state: Current optimization state with object positions and
-                optional per-env bounding boxes (accessed via ``state.get_bbox``).
+                optional per-env bounding boxes (accessed via state.get_bbox).
             debug: If True, print detailed loss breakdown.
 
         Returns:
@@ -233,10 +233,10 @@ def solve(
             initial_positions: List of dicts (one per env). Use a single-element list
                 for single-env placement.
             env_bboxes: Optional per-env bounding boxes keyed by object.
-                When provided, each ``AxisAlignedBoundingBox`` has shape
-                ``(batch, 3)`` so different batch rows can use different
-                geometry (heterogeneous placement). If ``None``, every row
-                uses the object's default ``get_bounding_box()``.
+                When provided, each AxisAlignedBoundingBox has shape
+                (batch, 3) so different batch rows can use different
+                geometry (heterogeneous placement). If None, every row
+                uses the object's default get_bounding_box().
 
         Returns:
             List of dicts (one per env) mapping objects to their solved (x, y, z) positions.
diff --git a/isaaclab_arena/relations/relation_solver_state.py b/isaaclab_arena/relations/relation_solver_state.py
index 3d837a377..fbb73219b 100644
--- a/isaaclab_arena/relations/relation_solver_state.py
+++ b/isaaclab_arena/relations/relation_solver_state.py
@@ -41,8 +41,8 @@ def __init__(
                 length > 1 = batched.
             device: Torch device for all tensors. Defaults to CPU.
             env_bboxes: Optional per-env bounding boxes keyed by object.
-                When provided, ``get_bbox(obj)`` returns the per-env bbox
-                (shape ``(batch, 3)``) instead of the object's default.
+                When provided, get_bbox(obj) returns the per-env bbox
+                (shape (batch, 3)) instead of the object's default.
         """
         assert len(initial_positions) >= 1, "initial_positions must contain at least one dict."
         anchor_objects = get_anchor_objects(objects)
@@ -150,10 +150,10 @@ def get_position(self, obj: ObjectBase) -> torch.Tensor:
         return self._optimizable_positions[:, opt_idx, :]
 
     def get_bbox(self, obj: ObjectBase) -> AxisAlignedBoundingBox:
-        """Return the local bounding box for *obj*, moved to the state's device.
+        """Return the local bounding box for obj, moved to the state's device.
 
-        Uses the per-env override from ``env_bboxes`` when available, otherwise
-        falls back to ``obj.get_bounding_box()``.
+        Uses the per-env override from env_bboxes when available, otherwise
+        falls back to obj.get_bounding_box().
         """
         if self._env_bboxes is not None and obj in self._env_bboxes:
             return self._env_bboxes[obj].to(self._device)

From 7cac0f15306b5cef0170492b0850732b8aa0dc81 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Sun, 17 May 2026 18:49:52 -0700
Subject: [PATCH 33/46] fix bug and change names

---
 isaaclab_arena/relations/bbox_helpers.py      | 42 ++++++-------------
 isaaclab_arena/relations/object_placer.py     |  4 +-
 .../relations/pooled_object_placer.py         | 30 ++++++-------
 isaaclab_arena/relations/relation_solver.py   |  4 ++
 .../tests/test_heterogeneous_placement.py     |  6 +--
 ...e_multi_object_no_collision_environment.py |  4 +-
 6 files changed, 39 insertions(+), 51 deletions(-)

diff --git a/isaaclab_arena/relations/bbox_helpers.py b/isaaclab_arena/relations/bbox_helpers.py
index b4a890aa2..aeeb33d16 100644
--- a/isaaclab_arena/relations/bbox_helpers.py
+++ b/isaaclab_arena/relations/bbox_helpers.py
@@ -3,14 +3,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Bounding-box utilities for the relation solver / object placer.
+"""Bounding-box helpers for heterogeneous placement.
 
-These helpers centralise the logic that expands single-object bounding boxes
-to per-environment tensors and detects whether any object carries
-env-specific geometry (e.g. RigidObjectSet with multiple variants).
-
-Placing this logic here (instead of on ObjectBase) keeps num_envs
-and other placement/simulation details out of the asset layer.
+Keeps num_envs and per-env geometry logic out of ObjectBase.
 """
 
 from __future__ import annotations
@@ -23,37 +18,24 @@
     from isaaclab_arena.assets.object_base import ObjectBase
 
 
-def object_has_env_specific_bboxes(obj: ObjectBase) -> bool:
-    """Whether obj produces different bounding boxes per environment.
-
-    Returns True for multi-variant RigidObjectSet instances (which
-    override get_bounding_box_per_env). All other objects return False.
-    """
+def is_heterogeneous(obj: ObjectBase) -> bool:
+    """True if obj has different bounding boxes per environment (e.g. multi-variant RigidObjectSet)."""
     return getattr(obj, "has_env_specific_bboxes", False)
 
 
-def any_object_has_env_specific_bboxes(objects: list[ObjectBase]) -> bool:
-    """Whether any object in the list has env-specific bounding boxes."""
-    return any(object_has_env_specific_bboxes(obj) for obj in objects)
+def has_heterogeneous_objects(objects: list[ObjectBase]) -> bool:
+    """True if any object in the list varies across environments."""
+    return any(is_heterogeneous(obj) for obj in objects)
 
 
 def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBoundingBox:
-    """Return per-environment local bounding boxes for obj.
-
-    For objects with env-specific variants (RigidObjectSet), delegates to the
-    object's own get_bounding_box_per_env override. For all other objects,
-    expands the single local bbox to (num_envs, 3).
-
-    Args:
-        obj: The object to query.
-        num_envs: Number of environments.
+    """Return bounding boxes expanded to (num_envs, 3).
 
-    Returns:
-        AxisAlignedBoundingBox with min_point / max_point of shape (num_envs, 3).
+    Heterogeneous objects delegate to their own get_bounding_box_per_env.
+    Homogeneous objects just broadcast the single bbox.
     """
-    override = getattr(obj, "get_bounding_box_per_env", None)
-    if override is not None:
-        return override(num_envs)
+    if is_heterogeneous(obj):
+        return obj.get_bounding_box_per_env(num_envs)
 
     bbox = obj.get_bounding_box()
     return AxisAlignedBoundingBox(
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index b90018f98..30858959e 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -10,7 +10,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-from isaaclab_arena.relations.bbox_helpers import any_object_has_env_specific_bboxes, get_bounding_box_per_env
+from isaaclab_arena.relations.bbox_helpers import get_bounding_box_per_env, has_heterogeneous_objects
 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
@@ -149,7 +149,7 @@ def place(
         #   - homogeneous: any solved layout serves any env; pick the top num_results.
         #   - heterogeneous: some objects vary per env, so each env owns a fixed slice
         #     of candidates and we pick the best within that slice.
-        uses_env_specific_bboxes = result_per_env and any_object_has_env_specific_bboxes(objects)
+        uses_env_specific_bboxes = result_per_env and has_heterogeneous_objects(objects)
 
         if uses_env_specific_bboxes:
             results_per_env = self._place_heterogeneous(
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 1595333df..a50b11e4c 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -9,7 +9,7 @@
 import torch
 from typing import TYPE_CHECKING
 
-from isaaclab_arena.relations.bbox_helpers import any_object_has_env_specific_bboxes, get_bounding_box_per_env
+from isaaclab_arena.relations.bbox_helpers import get_bounding_box_per_env, has_heterogeneous_objects
 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, PlacementResult
@@ -52,7 +52,7 @@ def __init__(
         # 1. Validate params.
         if pool_size < 1:
             raise ValueError(f"pool_size must be >= 1, got {pool_size}")
-        self._uses_env_specific_bboxes = any_object_has_env_specific_bboxes(objects)
+        self._uses_env_specific_bboxes = has_heterogeneous_objects(objects)
         if self._uses_env_specific_bboxes:
             assert num_envs is not None, "num_envs is required when layouts use env-specific object variants."
         self._num_envs = num_envs if num_envs is not None else 1
@@ -322,20 +322,22 @@ def had_fallbacks(self) -> bool:
         return self._had_fallbacks
 
     def sample_with_replacement(self, count: int) -> list[PlacementResult]:
-        """Pick count layouts at random per env-slot (non-consuming).
+        """Pick count layouts at random with replacement (non-consuming).
 
-        Slot i is filled by a random pick from env i % num_envs's
-        pool, so a length-count request walks env slots in order. Used
-        by resolve_on_reset=False to assign initial positions that persist
-        across resets.
+        For env-specific layouts, slot i picks from env i % num_envs's pool
+        so each result matches its absolute env. For reusable layouts, draws
+        are uniform IID from the full pool (preserving pre-heterogeneous behavior).
         """
-        results: list[PlacementResult] = []
-        for i in range(count):
-            cur_env = i % self._num_envs
-            pool = self._layout_pools[cur_env]
-            assert pool, f"Env {cur_env} has no valid layouts to sample from."
-            results.append(random.choice(pool))
-        return results
+        if self._uses_env_specific_bboxes:
+            results: list[PlacementResult] = []
+            for i in range(count):
+                cur_env = i % self._num_envs
+                pool = self._layout_pools[cur_env]
+                assert pool, f"Env {cur_env} has no valid layouts to sample from."
+                results.append(random.choice(pool))
+            return results
+        all_layouts = [layout for pool in self._layout_pools.values() for layout in pool]
+        return random.choices(all_layouts, k=count)
 
     @property
     def remaining(self) -> int:
diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index 72d037ce2..ef4013261 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -158,6 +158,10 @@ def _compute_no_overlap_loss(
 
         non_anchor_objects = state.optimizable_objects
         anchor_objects = list(state.anchor_objects)
+
+        # Skip no-overlap for On pairs: the On loss already pushes the child
+        # onto the parent surface, so penalizing bbox overlap between them
+        # would fight that constraint and cause oscillation.
         on_pairs: set[tuple[int, int]] = set()
         for obj in [*non_anchor_objects, *anchor_objects]:
             for rel in obj.get_relations():
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 07abb1173..0b0b838e0 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -14,7 +14,7 @@
 import pytest
 
 from isaaclab_arena.assets.dummy_object import DummyObject
-from isaaclab_arena.relations.bbox_helpers import get_bounding_box_per_env, object_has_env_specific_bboxes
+from isaaclab_arena.relations.bbox_helpers import get_bounding_box_per_env, is_heterogeneous
 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, PlacementResult
@@ -105,7 +105,7 @@ def test_dummy_object_preserves_constructor_relations():
     )
 
     assert obj.get_relations() == [anchor_relation]
-    assert object_has_env_specific_bboxes(obj) is False
+    assert is_heterogeneous(obj) is False
 
 
 def test_object_preserves_constructor_relations():
@@ -123,7 +123,7 @@ def test_object_preserves_constructor_relations():
     )
 
     assert obj.get_relations() == [anchor_relation]
-    assert object_has_env_specific_bboxes(obj) is False
+    assert is_heterogeneous(obj) is False
 
 
 # ---------------------------------------------------------------------------
diff --git a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
index e8bf8ab87..c40d1ec97 100644
--- a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
+++ b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
@@ -53,7 +53,7 @@
 
 # -- Heterogeneous mode default object sets ----------------------------------
 # Each entry is a multi-variant RigidObjectSet — each env gets a different
-# variant. Objects sourced from het-viz branch gif capture script.
+# variant.
 HETERO_VARIANT_SETS = {
     "bottles": [
         "mustard_bottle_hope_robolab",
@@ -199,7 +199,7 @@ def _build_heterogeneous_objects(self, tabletop_reference, object_names=None):
 
         When --objects is provided, each object becomes a single-variant RigidObjectSet.
         Otherwise, uses HETERO_FIXED_OBJECTS (pinned fruits) + HETERO_VARIANT_SETS
-        (multi-variant sets from het-viz branch).
+        (multi-variant sets).
         """
         from isaaclab_arena.assets.object_set import RigidObjectSet
         from isaaclab_arena.relations.relations import AtPosition, On

From 960fe9479bd54f8adba0008a6c32e3cbbb18d1f1 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Sun, 17 May 2026 19:09:44 -0700
Subject: [PATCH 34/46] revert redundant changes

---
 isaaclab_arena/assets/dummy_object.py                | 4 ++--
 isaaclab_arena/assets/object.py                      | 4 ++--
 isaaclab_arena/relations/relation_loss_strategies.py | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/isaaclab_arena/assets/dummy_object.py b/isaaclab_arena/assets/dummy_object.py
index cb7ab52a3..aabe95bfe 100644
--- a/isaaclab_arena/assets/dummy_object.py
+++ b/isaaclab_arena/assets/dummy_object.py
@@ -19,14 +19,14 @@ def __init__(
         name: str,
         bounding_box: AxisAlignedBoundingBox,
         initial_pose: Pose | None = None,
-        relations: list[RelationBase] | None = None,
+        relations: list[RelationBase] = [],
         **kwargs,
     ):
         self.name = name
         self.initial_pose = initial_pose
         self.bounding_box = bounding_box
         assert self.bounding_box is not None
-        self.relations = list(relations or [])
+        self.relations = list(relations)
 
     def add_relation(self, relation: RelationBase) -> None:
         self.relations.append(relation)
diff --git a/isaaclab_arena/assets/object.py b/isaaclab_arena/assets/object.py
index 6123e9f89..e2a92af36 100644
--- a/isaaclab_arena/assets/object.py
+++ b/isaaclab_arena/assets/object.py
@@ -32,7 +32,7 @@ def __init__(
         usd_path: str | None = None,
         scale: tuple[float, float, float] = (1.0, 1.0, 1.0),
         initial_pose: Pose | None = None,
-        relations: list[RelationBase] | None = None,
+        relations: list[RelationBase] = [],
         spawner_cfg: SpawnerCfg | None = None,
         **kwargs,
     ):
@@ -55,7 +55,7 @@ def __init__(
         self.spawner_cfg = spawner_cfg
         self.scale = scale
         self.initial_pose = initial_pose
-        self.relations = list(relations or [])
+        self.relations = list(relations)
         self.reset_pose = True
         self.spawn_cfg_addon = spawn_cfg_addon
         self.asset_cfg_addon = asset_cfg_addon
diff --git a/isaaclab_arena/relations/relation_loss_strategies.py b/isaaclab_arena/relations/relation_loss_strategies.py
index 6b0d08700..af42a2e98 100644
--- a/isaaclab_arena/relations/relation_loss_strategies.py
+++ b/isaaclab_arena/relations/relation_loss_strategies.py
@@ -18,7 +18,7 @@
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
-    from isaaclab_arena.relations.relations import AtPosition, NextTo, On, PositionLimits, Relation, UnaryRelation
+    from isaaclab_arena.relations.relations import AtPosition, NextTo, On, PositionLimits, Relation
 
 from isaaclab_arena.relations.relations import Side
 
@@ -71,7 +71,7 @@ class UnaryRelationLossStrategy(ABC):
     @abstractmethod
     def compute_loss(
         self,
-        relation: "UnaryRelation",
+        relation: "Relation",
         child_pos: torch.Tensor,
         child_bbox: AxisAlignedBoundingBox,
     ) -> torch.Tensor:

From 85395a565a09fbf8b540529ab1050f14e03b7f06 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 20 May 2026 13:59:06 -0700
Subject: [PATCH 35/46] address comments

---
 isaaclab_arena/assets/object_set.py           | 112 ++++++++++--------
 isaaclab_arena/relations/bbox_helpers.py      |  24 ++--
 isaaclab_arena/relations/object_placer.py     |   7 +-
 .../relations/pooled_object_placer.py         |   7 +-
 .../tests/test_heterogeneous_placement.py     |  43 ++++++-
 isaaclab_arena/tests/test_object_set.py       |   6 +-
 ...e_multi_object_no_collision_environment.py |  11 +-
 7 files changed, 133 insertions(+), 77 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 7bd2acb49..6add73a4d 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -48,6 +48,8 @@ def __init__(
             variant_indices_by_env: Optional fixed variant index for each environment.
             initial_pose: The initial pose of the object from this object set.
         """
+        if len(objects) < 2:
+            raise ValueError(f"Object set {name} must contain at least 2 objects.")
         if not self._are_all_objects_type_rigid(objects):
             raise ValueError(f"Object set {name} must contain only rigid objects.")
 
@@ -65,24 +67,21 @@ def __init__(
                     "rigid body at the same depth so paths match after rename. "
                     f"Rigid body depths by asset: {per_asset}."
                 )
-            self.object_usd_paths = self._modify_assets(objects)
-            print(f"Modified object USD paths: {self.object_usd_paths}")
+            self.member_usd_paths: list[str] = self._modify_assets(objects)
+            print(f"Modified object USD paths: {self.member_usd_paths}")
         else:
-            self.object_usd_paths = []
+            self.member_usd_paths = []
             for obj in objects:
                 assert obj.usd_path is not None
-                self.object_usd_paths.append(obj.usd_path)
+                self.member_usd_paths.append(obj.usd_path)
 
         self.objects: list[Object] = objects
-        self._member_object_usd_paths: list[str] = list(self.object_usd_paths)
         self.random_choice = random_choice
         self.variant_indices_by_env: list[int] | None = None
-        self.has_env_specific_bboxes = len(objects) > 1
 
         if variant_indices_by_env is not None:
             self._set_variant_indices_by_env(variant_indices_by_env)
 
-        # Set default prim_path if not provided
         if prim_path is None:
             prim_path = f"{{ENV_REGEX_NS}}/{name}"
 
@@ -91,11 +90,22 @@ def __init__(
             object_type=ObjectType.RIGID,
             usd_path="",
             prim_path=prim_path,
-            scale=(1.0, 1.0, 1.0),  # We rewrite the USDs to handle scaling
+            scale=(1.0, 1.0, 1.0),
             initial_pose=initial_pose,
             **kwargs,
         )
 
+    @property
+    def object_usd_paths(self) -> list[str]:
+        """Per-env USD paths passed to MultiUsdFileCfg.
+
+        Before variant indices are assigned, returns the member list.
+        After assignment, returns one path per env based on variant_indices_by_env.
+        """
+        if self.variant_indices_by_env is not None:
+            return [self.member_usd_paths[idx] for idx in self.variant_indices_by_env]
+        return self.member_usd_paths
+
     def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get the bounding box of the object set.
 
@@ -104,29 +114,30 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """
         return max(self.objects, key=lambda obj: obj.get_bounding_box().size[0, 2].item()).get_bounding_box()
 
-    def get_variant_indices(self, num_envs: int) -> list[int]:
-        """Return which member object index is assigned to each environment.
+    def assign_variants(self, num_envs: int) -> None:
+        """Assign one member-variant index per environment.
 
-        Multi-variant sets use one fixed assignment for the lifetime of the
-        object set. When ``random_choice`` is True, each env independently
-        samples one variant once. Otherwise, assignments repeat the member
-        order across environments.
+        The assignment is fixed for the lifetime of the object set: subsequent
+        calls with the same ``num_envs`` are no-ops, and a call with a
+        different ``num_envs`` raises. When ``random_choice`` is True, each
+        env independently samples one variant; otherwise assignments repeat
+        the member order across environments.
 
-        Args:
-            num_envs: Number of environments.
+        Callers (typically the placer once ``num_envs`` is known) must invoke
+        this before reading ``variant_indices_by_env`` or
+        ``get_bounding_box_per_env``.
 
-        Returns:
-            List of length ``num_envs`` with indices into ``self.objects``.
+        Args:
+            num_envs: Number of environments to assign variants for.
         """
-        if self.variant_indices_by_env is None:
-            self._set_variant_indices_by_env(self._generate_variant_indices(num_envs))
-        elif len(self.variant_indices_by_env) != num_envs:
-            raise ValueError(
-                f"RigidObjectSet '{self.name}' has variant assignments for "
-                f"{len(self.variant_indices_by_env)} envs, got request for {num_envs}."
-            )
-        assert self.variant_indices_by_env is not None
-        return self.variant_indices_by_env
+        if self.variant_indices_by_env is not None:
+            if len(self.variant_indices_by_env) != num_envs:
+                raise ValueError(
+                    f"RigidObjectSet '{self.name}' has variant assignments for "
+                    f"{len(self.variant_indices_by_env)} envs, got request for {num_envs}."
+                )
+            return
+        self._set_variant_indices_by_env(self._generate_variant_indices(num_envs))
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         """Get the actual bounding box for each env's variant.
@@ -135,52 +146,52 @@ def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         returns the real local bbox of the variant assigned to each env,
         enabling correct collision-free placement for heterogeneous scenes.
 
+        Requires ``assign_variants(num_envs)`` to have been called first.
+
         Args:
-            num_envs: Number of environments.
+            num_envs: Number of environments. Must match the assignment.
 
         Returns:
             ``AxisAlignedBoundingBox`` with ``min_point`` / ``max_point`` of
             shape ``(num_envs, 3)``.
         """
-        variant_indices = self.get_variant_indices(num_envs)
-        member_bboxes = [obj.get_bounding_box() for obj in self.objects]
+        assert self.variant_indices_by_env is not None, (
+            f"RigidObjectSet '{self.name}' has no variant assignment; "
+            "call assign_variants(num_envs) before get_bounding_box_per_env()."
+        )
+        assert len(self.variant_indices_by_env) == num_envs, (
+            f"RigidObjectSet '{self.name}' assigned for "
+            f"{len(self.variant_indices_by_env)} envs, got request for {num_envs}."
+        )
+        bounding_boxes = [obj.get_bounding_box() for obj in self.objects]
 
-        min_pts = torch.stack([member_bboxes[idx].min_point[0] for idx in variant_indices], dim=0)
-        max_pts = torch.stack([member_bboxes[idx].max_point[0] for idx in variant_indices], dim=0)
+        min_pts = torch.stack([bounding_boxes[idx].min_point[0] for idx in self.variant_indices_by_env], dim=0)
+        max_pts = torch.stack([bounding_boxes[idx].max_point[0] for idx in self.variant_indices_by_env], dim=0)
         return AxisAlignedBoundingBox(min_point=min_pts, max_point=max_pts)
 
     def get_contact_sensor_cfg(self, contact_against_object: ObjectBase | None = None) -> ContactSensorCfg:
         # We assume that by here, our USDs have been modified to be compatible with each other
         # and we can use the canonical first member USD to find the shallowest rigid body.
-        return super().get_contact_sensor_cfg(contact_against_object, usd_path=self._member_object_usd_paths[0])
+        return super().get_contact_sensor_cfg(contact_against_object, usd_path=self.member_usd_paths[0])
 
     def _generate_variant_indices(self, num_envs: int) -> list[int]:
         n = len(self.objects)
-        if n == 1:
-            return [0 for _ in range(num_envs)]
         if not self.random_choice:
             return [env_idx % n for env_idx in range(num_envs)]
         return torch.randint(low=0, high=n, size=(num_envs,)).tolist()
 
     def _set_variant_indices_by_env(self, variant_indices_by_env: list[int]) -> None:
+        """Validate and store variant indices, then sync the spawn config's USD path list."""
         n = len(self.objects)
-        if any(idx < 0 or idx >= n for idx in variant_indices_by_env):
-            raise ValueError(
-                f"RigidObjectSet '{self.name}' variant indices must be in [0, {n}); got {variant_indices_by_env}."
-            )
-
+        assert all(
+            0 <= idx < n for idx in variant_indices_by_env
+        ), f"RigidObjectSet '{self.name}' variant indices must be in [0, {n}); got {variant_indices_by_env}."
         self.variant_indices_by_env = list(variant_indices_by_env)
-        if len(self.objects) > 1:
-            # Keep Isaac Lab's MultiUsdFileCfg aligned with the fixed per-env variant assignment.
-            self.object_usd_paths = [self._member_object_usd_paths[idx] for idx in self.variant_indices_by_env]
-            spawn_cfg = self.object_cfg.spawn if getattr(self, "object_cfg", None) is not None else None
-            if isinstance(spawn_cfg, sim_utils.MultiUsdFileCfg):
-                spawn_cfg.usd_path = self.object_usd_paths
-                spawn_cfg.random_choice = False
+        spawn_cfg = self.object_cfg.spawn if getattr(self, "object_cfg", None) is not None else None
+        if isinstance(spawn_cfg, sim_utils.MultiUsdFileCfg):
+            spawn_cfg.usd_path = self.object_usd_paths
 
     def _are_all_objects_type_rigid(self, objects: list[Object]) -> bool:
-        if objects is None or len(objects) == 0:
-            raise ValueError(f"Object set {self.name} must contain at least 1 object.")
         for obj in objects:
             assert obj.usd_path is not None
             if detect_object_type(usd_path=obj.usd_path) != ObjectType.RIGID:
@@ -193,12 +204,11 @@ def _generate_rigid_cfg(self) -> RigidObjectCfg:
             prim_path=self.prim_path,
             spawn=sim_utils.MultiUsdFileCfg(
                 usd_path=self.object_usd_paths,
-                random_choice=self.random_choice if self.variant_indices_by_env is None else False,
+                random_choice=False,
                 activate_contact_sensors=True,
             ),
         )
         object_cfg = self._add_initial_pose_to_cfg(object_cfg)
-        assert isinstance(object_cfg, RigidObjectCfg)
         return object_cfg
 
     def _generate_articulation_cfg(self):
diff --git a/isaaclab_arena/relations/bbox_helpers.py b/isaaclab_arena/relations/bbox_helpers.py
index aeeb33d16..a44d2029a 100644
--- a/isaaclab_arena/relations/bbox_helpers.py
+++ b/isaaclab_arena/relations/bbox_helpers.py
@@ -12,29 +12,35 @@
 
 from typing import TYPE_CHECKING
 
+from isaaclab_arena.assets.object_set import RigidObjectSet
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
     from isaaclab_arena.assets.object_base import ObjectBase
 
 
-def is_heterogeneous(obj: ObjectBase) -> bool:
-    """True if obj has different bounding boxes per environment (e.g. multi-variant RigidObjectSet)."""
-    return getattr(obj, "has_env_specific_bboxes", False)
+def has_heterogeneous_objects(objects: list[ObjectBase]) -> bool:
+    """True if any object in the list is a RigidObjectSet."""
+    return any(isinstance(obj, RigidObjectSet) for obj in objects)
 
 
-def has_heterogeneous_objects(objects: list[ObjectBase]) -> bool:
-    """True if any object in the list varies across environments."""
-    return any(is_heterogeneous(obj) for obj in objects)
+def assign_variants_for_envs(objects: list[ObjectBase], num_envs: int) -> None:
+    """Assign per-env variants on every RigidObjectSet in the list.
+
+    Placers call this at the boundary before any per-env geometry reads.
+    """
+    for obj in objects:
+        if isinstance(obj, RigidObjectSet):
+            obj.assign_variants(num_envs)
 
 
 def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBoundingBox:
     """Return bounding boxes expanded to (num_envs, 3).
 
-    Heterogeneous objects delegate to their own get_bounding_box_per_env.
-    Homogeneous objects just broadcast the single bbox.
+    RigidObjectSet delegates to its own get_bounding_box_per_env.
+    All other objects broadcast their single bbox.
     """
-    if is_heterogeneous(obj):
+    if isinstance(obj, RigidObjectSet):
         return obj.get_bounding_box_per_env(num_envs)
 
     bbox = obj.get_bounding_box()
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index 30858959e..14a91d36b 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -10,7 +10,11 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-from isaaclab_arena.relations.bbox_helpers import get_bounding_box_per_env, has_heterogeneous_objects
+from isaaclab_arena.relations.bbox_helpers import (
+    assign_variants_for_envs,
+    get_bounding_box_per_env,
+    has_heterogeneous_objects,
+)
 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
@@ -254,6 +258,7 @@ def _place_heterogeneous(
                 (num_envs, 3). When None we call
                 get_bounding_box_per_env(obj, num_envs) on each object.
         """
+        assign_variants_for_envs(objects, num_envs)
         if env_bboxes is None:
             env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
 
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index a50b11e4c..fba894c74 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -9,7 +9,11 @@
 import torch
 from typing import TYPE_CHECKING
 
-from isaaclab_arena.relations.bbox_helpers import get_bounding_box_per_env, has_heterogeneous_objects
+from isaaclab_arena.relations.bbox_helpers import (
+    assign_variants_for_envs,
+    get_bounding_box_per_env,
+    has_heterogeneous_objects,
+)
 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, PlacementResult
@@ -182,6 +186,7 @@ def _solve_layouts_with_env_bboxes(self, num_layouts: int) -> tuple[list[Placeme
         layouts_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
         total_layouts = layouts_per_env * self._num_envs
 
+        assign_variants_for_envs(self._objects, self._num_envs)
         real_bboxes = {obj: get_bounding_box_per_env(obj, self._num_envs) for obj in self._objects}
 
         # (num_envs, 3) -> repeat each env's row layouts_per_env times -> (total_layouts, 3).
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 0b0b838e0..8b878376c 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -14,7 +14,8 @@
 import pytest
 
 from isaaclab_arena.assets.dummy_object import DummyObject
-from isaaclab_arena.relations.bbox_helpers import get_bounding_box_per_env, is_heterogeneous
+from isaaclab_arena.assets.object_set import RigidObjectSet
+from isaaclab_arena.relations.bbox_helpers import get_bounding_box_per_env
 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, PlacementResult
@@ -25,6 +26,39 @@
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 from isaaclab_arena.utils.pose import Pose
 
+# ---------------------------------------------------------------------------
+# Fixture: let HeterogeneousDummyObject trigger the heterogeneous path
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(autouse=True)
+def _patch_bbox_helpers_for_test_doubles(monkeypatch):
+    """Allow HeterogeneousDummyObject to trigger the heterogeneous placement path."""
+    from isaaclab_arena.relations import bbox_helpers
+
+    _original_has_het = bbox_helpers.has_heterogeneous_objects
+    _original_get_bbox = bbox_helpers.get_bounding_box_per_env
+
+    def _patched_has_het(objects):
+        if _original_has_het(objects):
+            return True
+        return any(hasattr(obj, "get_bounding_box_per_env") for obj in objects)
+
+    def _patched_get_bbox(obj, num_envs):
+        if isinstance(obj, RigidObjectSet):
+            return _original_get_bbox(obj, num_envs)
+        if hasattr(obj, "get_bounding_box_per_env"):
+            return obj.get_bounding_box_per_env(num_envs)
+        return _original_get_bbox(obj, num_envs)
+
+    monkeypatch.setattr("isaaclab_arena.relations.bbox_helpers.has_heterogeneous_objects", _patched_has_het)
+    monkeypatch.setattr("isaaclab_arena.relations.bbox_helpers.get_bounding_box_per_env", _patched_get_bbox)
+    monkeypatch.setattr("isaaclab_arena.relations.object_placer.has_heterogeneous_objects", _patched_has_het)
+    monkeypatch.setattr("isaaclab_arena.relations.object_placer.get_bounding_box_per_env", _patched_get_bbox)
+    monkeypatch.setattr("isaaclab_arena.relations.pooled_object_placer.has_heterogeneous_objects", _patched_has_het)
+    monkeypatch.setattr("isaaclab_arena.relations.pooled_object_placer.get_bounding_box_per_env", _patched_get_bbox)
+
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -40,7 +74,6 @@ class HeterogeneousDummyObject(DummyObject):
     def __init__(self, name: str, bboxes: list[AxisAlignedBoundingBox], **kwargs):
         super().__init__(name=name, bounding_box=bboxes[0], **kwargs)
         self._per_env_bboxes = bboxes
-        self.has_env_specific_bboxes = True
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         """Return env-specific bbox variants for this test double."""
@@ -87,7 +120,7 @@ def test_heterogeneous_dummy_returns_different_bboxes():
     large = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.3))
     obj = HeterogeneousDummyObject(name="set", bboxes=[small, large])
 
-    per_env = get_bounding_box_per_env(obj, 4)
+    per_env = obj.get_bounding_box_per_env(4)
     assert per_env.max_point.shape == (4, 3)
     # env 0 and 2 should use small; env 1 and 3 should use large
     assert torch.allclose(per_env.max_point[0], torch.tensor([0.1, 0.1, 0.1]))
@@ -105,7 +138,7 @@ def test_dummy_object_preserves_constructor_relations():
     )
 
     assert obj.get_relations() == [anchor_relation]
-    assert is_heterogeneous(obj) is False
+    assert not isinstance(obj, RigidObjectSet)
 
 
 def test_object_preserves_constructor_relations():
@@ -123,7 +156,7 @@ def test_object_preserves_constructor_relations():
     )
 
     assert obj.get_relations() == [anchor_relation]
-    assert is_heterogeneous(obj) is False
+    assert not isinstance(obj, RigidObjectSet)
 
 
 # ---------------------------------------------------------------------------
diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py
index b32dc3e27..4a464365c 100644
--- a/isaaclab_arena/tests/test_object_set.py
+++ b/isaaclab_arena/tests/test_object_set.py
@@ -48,7 +48,8 @@ def _test_object_set_samples_and_stores_variant_indices(simulation_app):
     ):
         obj_set = RigidObjectSet(name="cans", objects=[can_a, can_b], random_choice=True)
         assert obj_set.variant_indices_by_env is None
-        assert obj_set.get_variant_indices(num_envs=4) == assigned_variant_indices
+        obj_set.assign_variants(num_envs=4)
+        assert obj_set.variant_indices_by_env == assigned_variant_indices
 
     assert obj_set.object_usd_paths == [can_b.usd_path, can_a.usd_path, can_b.usd_path, can_b.usd_path]
     spawn_cfg = obj_set.object_cfg.spawn
@@ -79,7 +80,8 @@ def _test_object_set_default_variant_indices_follow_member_order(simulation_app)
         patch("isaaclab_arena.assets.object_set.find_shallowest_rigid_body", return_value="/rigid"),
     ):
         obj_set = RigidObjectSet(name="ordered_cans", objects=[can_a, can_b])
-        assert obj_set.get_variant_indices(num_envs=5) == [0, 1, 0, 1, 0]
+        obj_set.assign_variants(num_envs=5)
+        assert obj_set.variant_indices_by_env == [0, 1, 0, 1, 0]
 
     assert obj_set.object_usd_paths == [can_a.usd_path, can_b.usd_path, can_a.usd_path, can_b.usd_path, can_a.usd_path]
     spawn_cfg = obj_set.object_cfg.spawn
diff --git a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
index c40d1ec97..658ac6807 100644
--- a/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
+++ b/isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py
@@ -197,7 +197,7 @@ def _build_homogeneous_objects(self, tabletop_reference, object_names=None):
     def _build_heterogeneous_objects(self, tabletop_reference, object_names=None):
         """Build placeable objects for heterogeneous mode.
 
-        When --objects is provided, each object becomes a single-variant RigidObjectSet.
+        When --objects is provided, each object is placed directly (no per-env variance).
         Otherwise, uses HETERO_FIXED_OBJECTS (pinned fruits) + HETERO_VARIANT_SETS
         (multi-variant sets).
         """
@@ -207,16 +207,11 @@ def _build_heterogeneous_objects(self, tabletop_reference, object_names=None):
         # TODO(@zhx06): Address residual object bouncing with xy-only no-collision
         # constraints and anchor constraint handling in a follow-up change.
         if object_names:
-            print(
-                "Warning: --objects with --mode heterogeneous wraps each object as a "
-                "single-variant set (no per-env variance). Use default sets for true heterogeneity."
-            )
             placeable_assets = []
             for name in object_names:
                 obj = self.asset_registry.get_asset_by_name(name)()
-                obj_set = RigidObjectSet(name=name, objects=[obj])
-                obj_set.add_relation(On(tabletop_reference, clearance_m=0.01))
-                placeable_assets.append(obj_set)
+                obj.add_relation(On(tabletop_reference, clearance_m=0.01))
+                placeable_assets.append(obj)
         else:
             placeable_assets = []
             for name, x, y in HETERO_FIXED_OBJECTS:

From 9e58bdde234e6953b2d2c72b4ef3f3d6893e4ca9 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Thu, 21 May 2026 13:31:25 -0700
Subject: [PATCH 36/46] address comments

---
 isaaclab_arena/assets/object_set.py           |  13 +-
 ...box_helpers.py => bounding_box_helpers.py} |   9 +-
 isaaclab_arena/relations/object_placer.py     | 169 ++++++++++--------
 .../relations/pooled_object_placer.py         |  54 ++----
 .../tests/test_heterogeneous_placement.py     |  50 ++++--
 5 files changed, 158 insertions(+), 137 deletions(-)
 rename isaaclab_arena/relations/{bbox_helpers.py => bounding_box_helpers.py} (82%)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 6add73a4d..5c6052805 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -90,7 +90,7 @@ def __init__(
             object_type=ObjectType.RIGID,
             usd_path="",
             prim_path=prim_path,
-            scale=(1.0, 1.0, 1.0),
+            scale=(1.0, 1.0, 1.0),  # We rewrite the USDs to handle scaling.
             initial_pose=initial_pose,
             **kwargs,
         )
@@ -109,8 +109,9 @@ def object_usd_paths(self) -> list[str]:
     def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get the bounding box of the object set.
 
-        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.
+        This compatibility fallback returns the member bbox with the greatest
+        z-extent. Heterogeneous placement should use get_bounding_box_per_env()
+        after assign_variants() so each env uses its actual variant geometry.
         """
         return max(self.objects, key=lambda obj: obj.get_bounding_box().size[0, 2].item()).get_bounding_box()
 
@@ -142,9 +143,9 @@ def assign_variants(self, num_envs: int) -> None:
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         """Get the actual bounding box for each env's variant.
 
-        Unlike ``get_bounding_box()`` (which uses a max-z heuristic), this
-        returns the real local bbox of the variant assigned to each env,
-        enabling correct collision-free placement for heterogeneous scenes.
+        Unlike the single-bbox compatibility fallback, this returns the real
+        local bbox of the variant assigned to each env, enabling correct
+        collision-free placement for heterogeneous scenes.
 
         Requires ``assign_variants(num_envs)`` to have been called first.
 
diff --git a/isaaclab_arena/relations/bbox_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py
similarity index 82%
rename from isaaclab_arena/relations/bbox_helpers.py
rename to isaaclab_arena/relations/bounding_box_helpers.py
index a44d2029a..c7dcef150 100644
--- a/isaaclab_arena/relations/bbox_helpers.py
+++ b/isaaclab_arena/relations/bounding_box_helpers.py
@@ -12,7 +12,6 @@
 
 from typing import TYPE_CHECKING
 
-from isaaclab_arena.assets.object_set import RigidObjectSet
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
@@ -21,6 +20,8 @@
 
 def has_heterogeneous_objects(objects: list[ObjectBase]) -> bool:
     """True if any object in the list is a RigidObjectSet."""
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+
     return any(isinstance(obj, RigidObjectSet) for obj in objects)
 
 
@@ -28,9 +29,10 @@ def assign_variants_for_envs(objects: list[ObjectBase], num_envs: int) -> None:
     """Assign per-env variants on every RigidObjectSet in the list.
 
     Placers call this at the boundary before any per-env geometry reads.
+    Only called after has_heterogeneous_objects confirms the list contains RigidObjectSets.
     """
     for obj in objects:
-        if isinstance(obj, RigidObjectSet):
+        if hasattr(obj, "assign_variants"):
             obj.assign_variants(num_envs)
 
 
@@ -39,8 +41,9 @@ def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBound
 
     RigidObjectSet delegates to its own get_bounding_box_per_env.
     All other objects broadcast their single bbox.
+    Only called after has_heterogeneous_objects confirms heterogeneous objects exist.
     """
-    if isinstance(obj, RigidObjectSet):
+    if hasattr(obj, "get_bounding_box_per_env"):
         return obj.get_bounding_box_per_env(num_envs)
 
     bbox = obj.get_bounding_box()
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index 14a91d36b..bdece6335 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -10,7 +10,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-from isaaclab_arena.relations.bbox_helpers import (
+from isaaclab_arena.relations.bounding_box_helpers import (
     assign_variants_for_envs,
     get_bounding_box_per_env,
     has_heterogeneous_objects,
@@ -94,7 +94,6 @@ def place(
         objects: list[ObjectBase],
         num_envs: int = 1,
         result_per_env: bool = True,
-        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> PlacementResult | MultiEnvPlacementResult:
         """Place objects according to their spatial relations.
 
@@ -104,47 +103,18 @@ def place(
             num_envs: Number of environments. 1 for single-env; > 1 for batched
                 placement (one layout per env).
             result_per_env: When True (default), each environment gets a distinct
-                layout. When False, a single best layout is solved and applied
-                identically to all environments.
-            env_bboxes: Pre-computed per-env bounding boxes (shape (num_envs, 3)
-                per object). When provided, _place_heterogeneous uses these
-                instead of calling get_bounding_box_per_env(obj, num_envs). This
-                allows the caller to tile real-env bboxes for pooled solving.
+                layout. When False, homogeneous objects use a single best layout
+                for all environments. Heterogeneous objects always use per-env
+                layouts so each env is solved with its assigned variant geometry.
 
         Returns:
-            PlacementResult when a single layout is produced (num_envs=1 or
-            result_per_env=False); MultiEnvPlacementResult otherwise.
+            PlacementResult when a single layout is produced; MultiEnvPlacementResult
+            when multiple per-env layouts are produced.
         """
-        # Validate all objects have at least one relation
-        for obj in objects:
-            assert obj.get_relations(), (
-                f"Object '{obj.name}' has no relations. All objects passed to place() must have "
-                "at least one relation (e.g., On(), NextTo(), or IsAnchor())."
-            )
-
-        # Find all anchor objects
-        anchor_objects = get_anchor_objects(objects)
-        assert len(anchor_objects) > 0, (
-            "No anchor object found. Mark at least one object with IsAnchor() to serve as a fixed reference. "
-            "Example: table.add_relation(IsAnchor())"
-        )
-
-        # Validate all anchors have initial_pose set
-        for anchor in anchor_objects:
-            assert anchor.get_initial_pose() is not None, (
-                f"Anchor object '{anchor.name}' must have an initial_pose set. "
-                "Call anchor_object.set_initial_pose(...) before placing."
-            )
-
-        anchor_objects_set = set(anchor_objects)
-
-        # Create a local RNG generator from the seed so placement is reproducible without
-        # affecting the global torch RNG state (e.g. Isaac Sim's internal random streams).
-        generator: torch.Generator | None = None
-        if self.params.placement_seed is not None:
-            generator = torch.Generator()
+        anchor_objects_set, generator = self._prepare_placement(objects)
 
-        num_results = num_envs if result_per_env else 1
+        uses_env_specific_bboxes = has_heterogeneous_objects(objects)
+        num_results = num_envs if result_per_env or uses_env_specific_bboxes else 1
         max_attempts = self.params.max_placement_attempts
         num_candidates = max_attempts * num_results
 
@@ -153,31 +123,85 @@ def place(
         #   - homogeneous: any solved layout serves any env; pick the top num_results.
         #   - heterogeneous: some objects vary per env, so each env owns a fixed slice
         #     of candidates and we pick the best within that slice.
-        uses_env_specific_bboxes = result_per_env and has_heterogeneous_objects(objects)
-
         if uses_env_specific_bboxes:
-            results_per_env = self._place_heterogeneous(
+            ranked_results_per_env = self._place_heterogeneous(
                 objects,
                 anchor_objects_set,
                 num_envs,
                 max_attempts,
-                num_candidates,
                 generator,
-                env_bboxes=env_bboxes,
             )
+            results_per_env = [ranked_results[0] for ranked_results in ranked_results_per_env]
         else:
             results_per_env = self._place_homogeneous(
                 objects, anchor_objects_set, num_results, max_attempts, num_candidates, generator
             )
 
-        final_per_env = [r.positions for r in results_per_env]
+        positions_per_env = [r.positions for r in results_per_env]
         if self.params.apply_positions_to_objects:
-            self._apply_positions(final_per_env, anchor_objects_set)
+            self._apply_positions(positions_per_env, anchor_objects_set)
 
         if num_results == 1:
             return results_per_env[0]
         return MultiEnvPlacementResult(results=results_per_env)
 
+    def place_ranked_per_env(
+        self,
+        objects: list[ObjectBase],
+        num_envs: int,
+        results_per_env: int,
+    ) -> list[list[PlacementResult]]:
+        """Return ranked placement candidates per env.
+
+        This is used by PooledObjectPlacer to fill env-specific pools. The
+        return value has shape ``(num_envs, <=results_per_env)``: each outer
+        list entry corresponds to a real env, and each inner list is sorted with
+        valid lower-loss layouts first.
+        """
+        assert results_per_env > 0, f"results_per_env must be positive, got {results_per_env}"
+        anchor_objects_set, generator = self._prepare_placement(objects)
+        max_attempts = self.params.max_placement_attempts
+        ranked_results_per_env = self._place_heterogeneous(
+            objects,
+            anchor_objects_set,
+            num_envs,
+            max_attempts * results_per_env,
+            generator,
+        )
+
+        if self.params.apply_positions_to_objects:
+            top_positions_per_env = [ranked_results[0].positions for ranked_results in ranked_results_per_env]
+            self._apply_positions(top_positions_per_env, anchor_objects_set)
+
+        return [ranked_results[:results_per_env] for ranked_results in ranked_results_per_env]
+
+    def _prepare_placement(
+        self,
+        objects: list[ObjectBase],
+    ) -> tuple[set[ObjectBase], torch.Generator | None]:
+        """Validate placement inputs and create the local RNG generator."""
+        for obj in objects:
+            assert obj.get_relations(), (
+                f"Object '{obj.name}' has no relations. All objects passed to place() must have "
+                "at least one relation (e.g., On(), NextTo(), or IsAnchor())."
+            )
+
+        anchor_objects = get_anchor_objects(objects)
+        assert len(anchor_objects) > 0, (
+            "No anchor object found. Mark at least one object with IsAnchor() to serve as a fixed reference. "
+            "Example: table.add_relation(IsAnchor())"
+        )
+        for anchor in anchor_objects:
+            assert anchor.get_initial_pose() is not None, (
+                f"Anchor object '{anchor.name}' must have an initial_pose set. "
+                "Call anchor_object.set_initial_pose(...) before placing."
+            )
+
+        generator: torch.Generator | None = None
+        if self.params.placement_seed is not None:
+            generator = torch.Generator()
+        return set(anchor_objects), generator
+
     # ------------------------------------------------------------------
     # Placement strategies
     # ------------------------------------------------------------------
@@ -208,9 +232,10 @@ def _place_homogeneous(
         all_positions = self._solver.solve(objects, initial_positions)
         assert self._solver.last_loss_per_env is not None
         all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
+        all_validations = [self._validate_placement(positions) for positions in all_positions]
 
         all_candidates = [
-            PlacementCandidate(all_losses[idx], all_positions[idx], self._validate_placement(all_positions[idx]))
+            PlacementCandidate(all_losses[idx], all_positions[idx], all_validations[idx])
             for idx in range(num_candidates)
         ]
         all_candidates.sort(key=lambda candidate: (not candidate.is_valid, candidate.loss))
@@ -241,26 +266,20 @@ def _place_heterogeneous(
         objects: list[ObjectBase],
         anchor_objects_set: set[ObjectBase],
         num_envs: int,
-        max_attempts: int,
-        num_candidates: int,
+        candidates_per_env: int,
         generator: torch.Generator | None,
-        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
-    ) -> list[PlacementResult]:
+    ) -> list[list[PlacementResult]]:
         """Per-env placement: each candidate is tied to its env's object variants.
 
-        Candidates [cur_env * max_attempts : (cur_env + 1) * max_attempts]
+        Candidates [cur_env * candidates_per_env : (cur_env + 1) * candidates_per_env]
         belong to cur_env and use that env's variant geometry. We solve all
-        num_envs * max_attempts candidates in one batched solver.solve()
-        call and pick the best candidate within each env's slice.
-
-        Args:
-            env_bboxes: Optional pre-tiled per-env bboxes of shape
-                (num_envs, 3). When None we call
-                get_bounding_box_per_env(obj, num_envs) on each object.
+        num_envs * candidates_per_env candidates in one batched solver.solve()
+        call, rank candidates within each env slice, and return a nested list of
+        PlacementResult objects with shape (num_envs, candidates_per_env).
         """
         assign_variants_for_envs(objects, num_envs)
-        if env_bboxes is None:
-            env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
+        env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
+        num_candidates = num_envs * candidates_per_env
 
         # Build the per-env (1, 3) bbox views once: used both for On-guided
         # initial-position sampling and for per-env validation below.
@@ -276,16 +295,16 @@ def _place_heterogeneous(
         ]
 
         # Per-candidate bboxes (num_candidates, 3): each env's row repeated
-        # max_attempts times so candidates for that env share its geometry.
+        # candidates_per_env times so candidates for that env share its geometry.
         candidate_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {}
         for obj, bbox in env_bboxes.items():
-            min_pt = bbox.min_point.repeat_interleave(max_attempts, dim=0)
-            max_pt = bbox.max_point.repeat_interleave(max_attempts, dim=0)
+            min_pt = bbox.min_point.repeat_interleave(candidates_per_env, dim=0)
+            max_pt = bbox.max_point.repeat_interleave(candidates_per_env, dim=0)
             candidate_bboxes[obj] = AxisAlignedBoundingBox(min_point=min_pt, max_point=max_pt)
 
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
         for candidate_idx in range(num_candidates):
-            cur_env = candidate_idx // max_attempts
+            cur_env = candidate_idx // candidates_per_env
             if generator is not None:
                 assert self.params.placement_seed is not None
                 generator.manual_seed(self.params.placement_seed + candidate_idx)
@@ -298,9 +317,9 @@ def _place_heterogeneous(
         assert self._solver.last_loss_per_env is not None
         all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
 
-        results: list[PlacementResult] = []
+        ranked_results_per_env: list[list[PlacementResult]] = []
         for cur_env in range(num_envs):
-            start = cur_env * max_attempts
+            start = cur_env * candidates_per_env
             env_bbox_overrides = per_env_bbox_overrides[cur_env]
             env_candidates = [
                 PlacementCandidate(
@@ -308,21 +327,25 @@ def _place_heterogeneous(
                     all_positions[start + idx],
                     self._validate_placement(all_positions[start + idx], bbox_overrides=env_bbox_overrides),
                 )
-                for idx in range(max_attempts)
+                for idx in range(candidates_per_env)
             ]
             env_candidates.sort(key=lambda candidate: (not candidate.is_valid, candidate.loss))
-            best = env_candidates[0]
-            results.append(
+            ranked_results_per_env.append([
                 PlacementResult(
-                    success=best.is_valid, positions=best.positions, final_loss=best.loss, attempts=max_attempts
+                    success=candidate.is_valid,
+                    positions=candidate.positions,
+                    final_loss=candidate.loss,
+                    attempts=candidates_per_env,
                 )
-            )
+                for candidate in env_candidates
+            ])
 
+        results = [ranked_results[0] for ranked_results in ranked_results_per_env]
         if self.params.verbose:
             n_valid = sum(1 for r in results if r.success)
             print(f"Solved {num_candidates} candidates in one batch (heterogeneous): {n_valid}/{num_envs} env(s) valid")
 
-        return results
+        return ranked_results_per_env
 
     def _generate_initial_positions(
         self,
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index fba894c74..9c7ca8957 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -9,15 +9,10 @@
 import torch
 from typing import TYPE_CHECKING
 
-from isaaclab_arena.relations.bbox_helpers import (
-    assign_variants_for_envs,
-    get_bounding_box_per_env,
-    has_heterogeneous_objects,
-)
+from isaaclab_arena.relations.bounding_box_helpers import has_heterogeneous_objects
 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, PlacementResult
-from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
     from isaaclab_arena.assets.object_base import ObjectBase
@@ -117,8 +112,8 @@ def _solve_and_store(self, num_layouts: int) -> None:
 
             batch_size = max_missing * self._num_envs
             if self._uses_env_specific_bboxes:
-                all_results, layouts_per_env = self._solve_layouts_with_env_bboxes(batch_size)
-                self._store_env_matched_results(all_results, layouts_per_env)
+                ranked_results_per_env, layouts_per_env = self._solve_env_ranked_layouts(batch_size)
+                self._store_env_matched_results(ranked_results_per_env, layouts_per_env)
             else:
                 layouts = self._solve_reusable_layouts(batch_size)
                 self._store_reusable_results(layouts)
@@ -175,41 +170,27 @@ def _store_reusable_results(self, layouts: list[PlacementResult]) -> None:
             self._layout_pools[cur_env].append(layout)
             available[cur_env] += 1
 
-    def _solve_layouts_with_env_bboxes(self, num_layouts: int) -> tuple[list[PlacementResult], int]:
-        """Solve layouts tied to each env's actual object geometry.
+    def _solve_env_ranked_layouts(self, num_layouts: int) -> tuple[list[list[PlacementResult]], int]:
+        """Solve ranked layouts tied to each env's actual object geometry.
 
-        Computes bounding boxes for the real num_envs once, tiles them
-        to num_layouts entries, and solves everything in one batched
-        place() call. Result i is mapped back to real env
-        i // layouts_per_env for pool storage.
+        Returns ranked candidate lists per real env so the pool can store
+        multiple layouts for each env without pretending extra candidate rows
+        are extra environments.
         """
         layouts_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
-        total_layouts = layouts_per_env * self._num_envs
-
-        assign_variants_for_envs(self._objects, self._num_envs)
-        real_bboxes = {obj: get_bounding_box_per_env(obj, self._num_envs) for obj in self._objects}
-
-        # (num_envs, 3) -> repeat each env's row layouts_per_env times -> (total_layouts, 3).
-        tiled_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {
-            obj: AxisAlignedBoundingBox(
-                min_point=bbox.min_point.repeat_interleave(layouts_per_env, dim=0),
-                max_point=bbox.max_point.repeat_interleave(layouts_per_env, dim=0),
-            )
-            for obj, bbox in real_bboxes.items()
-        }
 
         with torch.inference_mode(False):
-            result = self._placer.place(
+            ranked_results_per_env = self._placer.place_ranked_per_env(
                 self._objects,
-                num_envs=total_layouts,
-                result_per_env=True,
-                env_bboxes=tiled_bboxes,
+                num_envs=self._num_envs,
+                results_per_env=layouts_per_env,
             )
 
-        all_results = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
-        return all_results, layouts_per_env
+        return ranked_results_per_env, layouts_per_env
 
-    def _store_env_matched_results(self, all_results: list[PlacementResult], layouts_per_env: int) -> None:
+    def _store_env_matched_results(
+        self, ranked_results_per_env: list[list[PlacementResult]], layouts_per_env: int
+    ) -> None:
         """Store env-matched results into their corresponding pools.
 
         Prefer successful layouts for each env. If a specific env has no valid
@@ -219,8 +200,7 @@ def _store_env_matched_results(self, all_results: list[PlacementResult], layouts
         total_valid = 0
         fallback_envs = []
         for cur_env in range(self._num_envs):
-            start = cur_env * layouts_per_env
-            env_results = all_results[start : start + layouts_per_env]
+            env_results = ranked_results_per_env[cur_env][:layouts_per_env]
             valid_results = [r for r in env_results if r.success]
             if valid_results:
                 self._layout_pools[cur_env].extend(valid_results)
@@ -230,7 +210,7 @@ def _store_env_matched_results(self, all_results: list[PlacementResult], layouts
                 fallback_envs.append(cur_env)
                 self._had_fallbacks = True
 
-        total_solved = len(all_results)
+        total_solved = sum(min(len(env_results), layouts_per_env) for env_results in ranked_results_per_env)
         if total_valid < total_solved:
             msg = (
                 f"Placement pool (env-specific bbox layouts): solved {total_solved} candidates,"
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 8b878376c..89c0eff89 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -14,8 +14,7 @@
 import pytest
 
 from isaaclab_arena.assets.dummy_object import DummyObject
-from isaaclab_arena.assets.object_set import RigidObjectSet
-from isaaclab_arena.relations.bbox_helpers import get_bounding_box_per_env
+from isaaclab_arena.relations.bounding_box_helpers import get_bounding_box_per_env
 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, PlacementResult
@@ -32,31 +31,25 @@
 
 
 @pytest.fixture(autouse=True)
-def _patch_bbox_helpers_for_test_doubles(monkeypatch):
-    """Allow HeterogeneousDummyObject to trigger the heterogeneous placement path."""
-    from isaaclab_arena.relations import bbox_helpers
+def _patch_bounding_box_helpers_for_test_doubles(monkeypatch):
+    """Allow HeterogeneousDummyObject to trigger the heterogeneous placement path.
 
-    _original_has_het = bbox_helpers.has_heterogeneous_objects
-    _original_get_bbox = bbox_helpers.get_bounding_box_per_env
+    Only has_heterogeneous_objects needs patching — it uses isinstance(obj, RigidObjectSet)
+    which won't match test doubles. The downstream functions (assign_variants_for_envs,
+    get_bounding_box_per_env) already duck-type via hasattr.
+    """
+    from isaaclab_arena.relations import bounding_box_helpers
+
+    _original_has_het = bounding_box_helpers.has_heterogeneous_objects
 
     def _patched_has_het(objects):
         if _original_has_het(objects):
             return True
         return any(hasattr(obj, "get_bounding_box_per_env") for obj in objects)
 
-    def _patched_get_bbox(obj, num_envs):
-        if isinstance(obj, RigidObjectSet):
-            return _original_get_bbox(obj, num_envs)
-        if hasattr(obj, "get_bounding_box_per_env"):
-            return obj.get_bounding_box_per_env(num_envs)
-        return _original_get_bbox(obj, num_envs)
-
-    monkeypatch.setattr("isaaclab_arena.relations.bbox_helpers.has_heterogeneous_objects", _patched_has_het)
-    monkeypatch.setattr("isaaclab_arena.relations.bbox_helpers.get_bounding_box_per_env", _patched_get_bbox)
+    monkeypatch.setattr("isaaclab_arena.relations.bounding_box_helpers.has_heterogeneous_objects", _patched_has_het)
     monkeypatch.setattr("isaaclab_arena.relations.object_placer.has_heterogeneous_objects", _patched_has_het)
-    monkeypatch.setattr("isaaclab_arena.relations.object_placer.get_bounding_box_per_env", _patched_get_bbox)
     monkeypatch.setattr("isaaclab_arena.relations.pooled_object_placer.has_heterogeneous_objects", _patched_has_het)
-    monkeypatch.setattr("isaaclab_arena.relations.pooled_object_placer.get_bounding_box_per_env", _patched_get_bbox)
 
 
 # ---------------------------------------------------------------------------
@@ -129,6 +122,7 @@ def test_heterogeneous_dummy_returns_different_bboxes():
 
 def test_dummy_object_preserves_constructor_relations():
     """DummyObject should keep relations passed at construction time."""
+    from isaaclab_arena.assets.object_set import RigidObjectSet
 
     anchor_relation = IsAnchor()
     obj = DummyObject(
@@ -145,6 +139,7 @@ def test_object_preserves_constructor_relations():
     """Object should keep relations passed at construction time."""
     from isaaclab_arena.assets.object import Object
     from isaaclab_arena.assets.object_base import ObjectType
+    from isaaclab_arena.assets.object_set import RigidObjectSet
 
     anchor_relation = IsAnchor()
     obj = Object(
@@ -320,6 +315,25 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
         ).item(), f"Env {env_idx}: A and X bboxes overlap at positions A={r.positions[obj_a]}, X={r.positions[obj_x]}"
 
 
+def test_heterogeneous_placement_always_returns_per_env_results():
+    """Heterogeneous placement should not fall back to one shared approximate layout."""
+
+    desk, hetero, _placer_params = _make_hetero_pool_objects()
+
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    params = ObjectPlacerParams(
+        solver_params=solver_params,
+        apply_positions_to_objects=False,
+        placement_seed=42,
+    )
+
+    placer = ObjectPlacer(params=params)
+    result = placer.place([desk, hetero], num_envs=4, result_per_env=False)
+
+    assert isinstance(result, MultiEnvPlacementResult)
+    assert len(result.results) == 4
+
+
 def test_object_placer_homogeneous_path_returns_multi_env_result():
     """When no heterogeneous objects exist, the homogeneous path is used."""
 

From fee37d23e582ecf12b07fe5f9f1566210a779a41 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Thu, 21 May 2026 17:50:19 -0700
Subject: [PATCH 37/46] address review comments

---
 isaaclab_arena/assets/object_set.py           |  9 +++--
 isaaclab_arena/relations/object_placer.py     |  7 ++--
 .../relations/pooled_object_placer.py         |  2 --
 .../tests/test_heterogeneous_placement.py     | 36 +++++++++++++++++--
 isaaclab_arena/tests/test_object_set.py       | 27 ++++++++++++++
 5 files changed, 69 insertions(+), 12 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 5c6052805..d183dfd87 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -132,11 +132,10 @@ def assign_variants(self, num_envs: int) -> None:
             num_envs: Number of environments to assign variants for.
         """
         if self.variant_indices_by_env is not None:
-            if len(self.variant_indices_by_env) != num_envs:
-                raise ValueError(
-                    f"RigidObjectSet '{self.name}' has variant assignments for "
-                    f"{len(self.variant_indices_by_env)} envs, got request for {num_envs}."
-                )
+            assert len(self.variant_indices_by_env) == num_envs, (
+                f"RigidObjectSet '{self.name}' has variant assignments for "
+                f"{len(self.variant_indices_by_env)} envs, got request for {num_envs}."
+            )
             return
         self._set_variant_indices_by_env(self._generate_variant_indices(num_envs))
 
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index bdece6335..ec3622cb7 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -104,8 +104,9 @@ def place(
                 placement (one layout per env).
             result_per_env: When True (default), each environment gets a distinct
                 layout. When False, homogeneous objects use a single best layout
-                for all environments. Heterogeneous objects always use per-env
-                layouts so each env is solved with its assigned variant geometry.
+                for all environments. Ignored for heterogeneous scenes; those
+                always produce one layout per env so each env is solved with its
+                assigned variant geometry.
 
         Returns:
             PlacementResult when a single layout is produced; MultiEnvPlacementResult
@@ -215,7 +216,7 @@ def _place_homogeneous(
         num_candidates: int,
         generator: torch.Generator | None,
     ) -> list[PlacementResult]:
-        """Original batched placement path.
+        """Batched placement path for objects with shared geometry.
 
         Solves max_attempts * num_results candidates in one batched
         solver.solve() call, then returns the best num_results ranked
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 9c7ca8957..a43d69e81 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -48,7 +48,6 @@ def __init__(
         pool_size: int = 100,
         num_envs: int | None = None,
     ) -> None:
-        # 1. Validate params.
         if pool_size < 1:
             raise ValueError(f"pool_size must be >= 1, got {pool_size}")
         self._uses_env_specific_bboxes = has_heterogeneous_objects(objects)
@@ -66,7 +65,6 @@ def __init__(
         self._layout_pools: dict[int, list[PlacementResult]] = {cur_env: [] for cur_env in range(self._num_envs)}
         self._layout_cursors: dict[int, int] = {cur_env: 0 for cur_env in range(self._num_envs)}
 
-        # 3. Solve the initial pool and assert every env has at least one layout.
         self._solve_and_store(pool_size)
         for cur_env, pool in self._layout_pools.items():
             if not pool:
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 89c0eff89..e0f786cbb 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -334,6 +334,31 @@ def test_heterogeneous_placement_always_returns_per_env_results():
     assert len(result.results) == 4
 
 
+def test_object_placer_place_ranked_per_env_returns_sorted_env_lists():
+    """place_ranked_per_env should return ranked candidate lists for each env."""
+
+    desk, hetero, _placer_params = _make_hetero_pool_objects()
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    params = ObjectPlacerParams(
+        solver_params=solver_params,
+        apply_positions_to_objects=False,
+        placement_seed=42,
+    )
+
+    placer = ObjectPlacer(params=params)
+    ranked_results = placer.place_ranked_per_env([desk, hetero], num_envs=3, results_per_env=2)
+
+    assert len(ranked_results) == 3
+    for env_results in ranked_results:
+        assert len(env_results) == 2
+        assert all(hetero in result.positions for result in env_results)
+        sort_keys = [(not result.success, result.final_loss) for result in env_results]
+        assert sort_keys == sorted(sort_keys)
+
+    with pytest.raises(AssertionError):
+        placer.place_ranked_per_env([desk, hetero], num_envs=3, results_per_env=0)
+
+
 def test_object_placer_homogeneous_path_returns_multi_env_result():
     """When no heterogeneous objects exist, the homogeneous path is used."""
 
@@ -413,9 +438,16 @@ def test_pooled_placer_heterogeneous_sample_with_replacement():
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
+    pool._layout_pools = {
+        env_id: [
+            PlacementResult(success=True, positions={hetero: (float(env_id), 0.0, 0.0)}, final_loss=0.0, attempts=1)
+        ]
+        for env_id in range(4)
+    }
+    pool._layout_cursors = {env_id: 0 for env_id in range(4)}
     initial_remaining = pool.remaining
-    samples = pool.sample_with_replacement(4)
-    assert len(samples) == 4
+    samples = pool.sample_with_replacement(8)
+    assert [sample.positions[hetero][0] for sample in samples] == [0.0, 1.0, 2.0, 3.0, 0.0, 1.0, 2.0, 3.0]
     assert pool.remaining == initial_remaining, "sample_with_replacement should not consume layouts"
 
 
diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py
index 4a464365c..b41b3f645 100644
--- a/isaaclab_arena/tests/test_object_set.py
+++ b/isaaclab_arena/tests/test_object_set.py
@@ -94,6 +94,25 @@ def _test_object_set_default_variant_indices_follow_member_order(simulation_app)
     return True
 
 
+def _test_object_set_rejects_variant_reassignment_with_different_num_envs(simulation_app):
+    """Variant assignment should remain fixed for the original env count."""
+    from isaaclab_arena.assets.object_base import ObjectType
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+
+    can_a, can_b, _bbox_a, _bbox_b = _make_object_set_variants()
+    with (
+        patch("isaaclab_arena.assets.object_set.detect_object_type", return_value=ObjectType.RIGID),
+        patch("isaaclab_arena.assets.object_set.find_shallowest_rigid_body", return_value="/rigid"),
+    ):
+        obj_set = RigidObjectSet(name="preassigned_cans", objects=[can_a, can_b], variant_indices_by_env=[0, 1, 0])
+
+    try:
+        obj_set.assign_variants(num_envs=4)
+    except AssertionError:
+        return True
+    return False
+
+
 def _build_and_reset_env(simulation_app, scene_assets, env_name="object_set_test", task=None):
     """Build arena env with given scene and optional task, then reset. Returns env (caller must close)."""
     from isaaclab_arena.assets.registries import AssetRegistry
@@ -454,6 +473,14 @@ def test_object_set_default_variant_indices_follow_member_order():
     assert result, f"Test {_test_object_set_default_variant_indices_follow_member_order.__name__} failed"
 
 
+def test_object_set_rejects_variant_reassignment_with_different_num_envs():
+    result = run_simulation_app_function(
+        _test_object_set_rejects_variant_reassignment_with_different_num_envs,
+        headless=HEADLESS,
+    )
+    assert result, f"Test {_test_object_set_rejects_variant_reassignment_with_different_num_envs.__name__} failed"
+
+
 def test_articulation_object_set():
     result = run_simulation_app_function(
         _test_articulation_object_set,

From f224b467b0bca96e0a70a4daac07286004fa526e Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Fri, 22 May 2026 10:11:39 -0700
Subject: [PATCH 38/46] alignment of homo and heter mode

---
 .../relations/bounding_box_helpers.py         |   3 +-
 isaaclab_arena/relations/object_placer.py     | 333 +++++++++---------
 .../relations/pooled_object_placer.py         |   5 +-
 .../tests/test_heterogeneous_placement.py     |  10 +
 4 files changed, 174 insertions(+), 177 deletions(-)

diff --git a/isaaclab_arena/relations/bounding_box_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py
index c7dcef150..82093ca14 100644
--- a/isaaclab_arena/relations/bounding_box_helpers.py
+++ b/isaaclab_arena/relations/bounding_box_helpers.py
@@ -29,7 +29,7 @@ def assign_variants_for_envs(objects: list[ObjectBase], num_envs: int) -> None:
     """Assign per-env variants on every RigidObjectSet in the list.
 
     Placers call this at the boundary before any per-env geometry reads.
-    Only called after has_heterogeneous_objects confirms the list contains RigidObjectSets.
+    Objects without variants are ignored.
     """
     for obj in objects:
         if hasattr(obj, "assign_variants"):
@@ -41,7 +41,6 @@ def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBound
 
     RigidObjectSet delegates to its own get_bounding_box_per_env.
     All other objects broadcast their single bbox.
-    Only called after has_heterogeneous_objects confirms heterogeneous objects exist.
     """
     if hasattr(obj, "get_bounding_box_per_env"):
         return obj.get_bounding_box_per_env(num_envs)
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index ec3622cb7..dfdb3ccc9 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -8,7 +8,7 @@
 import math
 import torch
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 from isaaclab_arena.relations.bounding_box_helpers import (
     assign_variants_for_envs,
@@ -34,7 +34,7 @@
 
 @dataclass
 class PlacementCandidate:
-    """A single solver result, ranked and selected in ObjectPlacer.place()."""
+    """A scored solver result used for ranking inside ObjectPlacer."""
 
     loss: float
     """Loss value returned by the solver."""
@@ -66,11 +66,6 @@ class ObjectPlacer:
     """
 
     def __init__(self, params: ObjectPlacerParams | None = None):
-        """Initialize the ObjectPlacer.
-
-        Args:
-            params: Configuration parameters. If None, uses defaults.
-        """
         self.params = params or ObjectPlacerParams()
         self._solver = RelationSolver(params=self.params.solver_params)
 
@@ -79,12 +74,7 @@ def _resolve_bbox(
         obj: ObjectBase,
         overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None,
     ) -> AxisAlignedBoundingBox:
-        """Return overrides[obj] if present, otherwise the object's default bbox.
-
-        Heterogeneous placement passes a single-env override bbox; the
-        homogeneous path and unit tests pass None and rely on the
-        object's own bbox.
-        """
+        """Return overrides[obj] if present, otherwise the object's default bbox."""
         if overrides is not None and obj in overrides:
             return overrides[obj]
         return obj.get_bounding_box()
@@ -115,33 +105,39 @@ def place(
         anchor_objects_set, generator = self._prepare_placement(objects)
 
         uses_env_specific_bboxes = has_heterogeneous_objects(objects)
+        # Env-specific variants cannot share one solved layout because each env
+        # must be solved against its assigned object geometry.
         num_results = num_envs if result_per_env or uses_env_specific_bboxes else 1
         max_attempts = self.params.max_placement_attempts
-        num_candidates = max_attempts * num_results
-
-        # Two solve paths produce a list[PlacementResult] of length num_results, sorted by
-        # (is_valid, loss):
-        #   - homogeneous: any solved layout serves any env; pick the top num_results.
-        #   - heterogeneous: some objects vary per env, so each env owns a fixed slice
-        #     of candidates and we pick the best within that slice.
+        # Solve max_attempts candidates per result so we can pick the best survivor.
         if uses_env_specific_bboxes:
-            ranked_results_per_env = self._place_heterogeneous(
+            ranked_results = self._place_ranked(
                 objects,
                 anchor_objects_set,
                 num_envs,
-                max_attempts,
-                generator,
+                candidates_per_env=max_attempts,
+                attempts_per_result=max_attempts,
+                ranking_mode="env_specific_geometry",
+                generator=generator,
             )
-            results_per_env = [ranked_results[0] for ranked_results in ranked_results_per_env]
+            results_per_env = [env_results[0] for env_results in ranked_results]
         else:
-            results_per_env = self._place_homogeneous(
-                objects, anchor_objects_set, num_results, max_attempts, num_candidates, generator
+            ranked_results = self._place_ranked(
+                objects,
+                anchor_objects_set,
+                num_envs=num_results,
+                candidates_per_env=max_attempts,
+                attempts_per_result=max_attempts,
+                ranking_mode="shared_geometry",
+                generator=generator,
             )
+            results_per_env = ranked_results[0][:num_results]
 
         positions_per_env = [r.positions for r in results_per_env]
         if self.params.apply_positions_to_objects:
             self._apply_positions(positions_per_env, anchor_objects_set)
 
+        # Single layout when callers don't need per-env results; otherwise return per-env layouts.
         if num_results == 1:
             return results_per_env[0]
         return MultiEnvPlacementResult(results=results_per_env)
@@ -154,33 +150,34 @@ def place_ranked_per_env(
     ) -> list[list[PlacementResult]]:
         """Return ranked placement candidates per env.
 
-        This is used by PooledObjectPlacer to fill env-specific pools. The
-        return value has shape ``(num_envs, <=results_per_env)``: each outer
-        list entry corresponds to a real env, and each inner list is sorted with
-        valid lower-loss layouts first.
+        Use this for PooledObjectPlacer, where each env pool needs multiple
+        candidate layouts and will apply poses later when sampled. Use place()
+        for the normal public API that returns the selected placement result.
+        The return value has shape ``(num_envs, <=results_per_env)``: each
+        outer list entry corresponds to a real env, and each inner list is
+        sorted with valid lower-loss layouts first.
         """
         assert results_per_env > 0, f"results_per_env must be positive, got {results_per_env}"
         anchor_objects_set, generator = self._prepare_placement(objects)
         max_attempts = self.params.max_placement_attempts
-        ranked_results_per_env = self._place_heterogeneous(
+        # Callers ask for independent candidate lists per env, not one shared-geometry top-N list.
+        ranked_results_per_env = self._place_ranked(
             objects,
             anchor_objects_set,
             num_envs,
-            max_attempts * results_per_env,
-            generator,
+            candidates_per_env=max_attempts * results_per_env,
+            attempts_per_result=max_attempts,
+            ranking_mode="env_specific_geometry",
+            generator=generator,
         )
 
-        if self.params.apply_positions_to_objects:
-            top_positions_per_env = [ranked_results[0].positions for ranked_results in ranked_results_per_env]
-            self._apply_positions(top_positions_per_env, anchor_objects_set)
-
         return [ranked_results[:results_per_env] for ranked_results in ranked_results_per_env]
 
     def _prepare_placement(
         self,
         objects: list[ObjectBase],
     ) -> tuple[set[ObjectBase], torch.Generator | None]:
-        """Validate placement inputs and create the local RNG generator."""
+        """Validate placement inputs and allocate an RNG seeded per candidate later."""
         for obj in objects:
             assert obj.get_relations(), (
                 f"Object '{obj.name}' has no relations. All objects passed to place() must have "
@@ -207,101 +204,34 @@ def _prepare_placement(
     # Placement strategies
     # ------------------------------------------------------------------
 
-    def _place_homogeneous(
-        self,
-        objects: list[ObjectBase],
-        anchor_objects_set: set[ObjectBase],
-        num_results: int,
-        max_attempts: int,
-        num_candidates: int,
-        generator: torch.Generator | None,
-    ) -> list[PlacementResult]:
-        """Batched placement path for objects with shared geometry.
-
-        Solves max_attempts * num_results candidates in one batched
-        solver.solve() call, then returns the best num_results ranked
-        by (is_valid, loss). All envs share the same geometry, so any solved
-        layout can serve any env.
-        """
-        initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
-        for candidate_idx in range(num_candidates):
-            if generator is not None:
-                assert self.params.placement_seed is not None
-                generator.manual_seed(self.params.placement_seed + candidate_idx)
-            initial_positions.append(self._generate_initial_positions(objects, anchor_objects_set, generator))
-
-        all_positions = self._solver.solve(objects, initial_positions)
-        assert self._solver.last_loss_per_env is not None
-        all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
-        all_validations = [self._validate_placement(positions) for positions in all_positions]
-
-        all_candidates = [
-            PlacementCandidate(all_losses[idx], all_positions[idx], all_validations[idx])
-            for idx in range(num_candidates)
-        ]
-        all_candidates.sort(key=lambda candidate: (not candidate.is_valid, candidate.loss))
-        selected = all_candidates[:num_results]
-
-        if self.params.verbose:
-            total_valid = sum(1 for candidate in all_candidates if candidate.is_valid)
-            finite_losses = [candidate.loss for candidate in all_candidates if math.isfinite(candidate.loss)]
-            mean_loss = sum(finite_losses) / len(finite_losses) if finite_losses else float("inf")
-            n_valid = sum(1 for candidate in selected if candidate.is_valid)
-            print(
-                f"Solved {num_candidates} candidates in one batch: mean loss = {mean_loss:.6f},"
-                f" {total_valid} valid, selected best {num_results} ({n_valid} valid)"
-            )
-
-        return [
-            PlacementResult(
-                success=candidate.is_valid,
-                positions=candidate.positions,
-                final_loss=candidate.loss,
-                attempts=max_attempts,
-            )
-            for candidate in selected
-        ]
-
-    def _place_heterogeneous(
+    def _place_ranked(
         self,
         objects: list[ObjectBase],
         anchor_objects_set: set[ObjectBase],
         num_envs: int,
         candidates_per_env: int,
+        attempts_per_result: int,
+        ranking_mode: Literal["shared_geometry", "env_specific_geometry"],
         generator: torch.Generator | None,
     ) -> list[list[PlacementResult]]:
-        """Per-env placement: each candidate is tied to its env's object variants.
+        """Solve and rank placement candidates.
 
-        Candidates [cur_env * candidates_per_env : (cur_env + 1) * candidates_per_env]
-        belong to cur_env and use that env's variant geometry. We solve all
-        num_envs * candidates_per_env candidates in one batched solver.solve()
-        call, rank candidates within each env slice, and return a nested list of
-        PlacementResult objects with shape (num_envs, candidates_per_env).
+        Shared-geometry ranking is used only when every candidate is
+        interchangeable. Env-specific-geometry ranking keeps candidates tied to
+        their env's assigned variants.
         """
+        if ranking_mode == "shared_geometry":
+            assert not has_heterogeneous_objects(objects), (
+                "Shared-geometry ranking requires homogeneous objects; candidates are not interchangeable across"
+                " variants."
+            )
+
         assign_variants_for_envs(objects, num_envs)
         env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
         num_candidates = num_envs * candidates_per_env
-
-        # Build the per-env (1, 3) bbox views once: used both for On-guided
-        # initial-position sampling and for per-env validation below.
-        per_env_bbox_overrides: list[dict[ObjectBase, AxisAlignedBoundingBox]] = [
-            {
-                obj: AxisAlignedBoundingBox(
-                    min_point=env_bboxes[obj].min_point[cur_env : cur_env + 1],
-                    max_point=env_bboxes[obj].max_point[cur_env : cur_env + 1],
-                )
-                for obj in objects
-            }
-            for cur_env in range(num_envs)
-        ]
-
-        # Per-candidate bboxes (num_candidates, 3): each env's row repeated
-        # candidates_per_env times so candidates for that env share its geometry.
-        candidate_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {}
-        for obj, bbox in env_bboxes.items():
-            min_pt = bbox.min_point.repeat_interleave(candidates_per_env, dim=0)
-            max_pt = bbox.max_point.repeat_interleave(candidates_per_env, dim=0)
-            candidate_bboxes[obj] = AxisAlignedBoundingBox(min_point=min_pt, max_point=max_pt)
+        candidate_bboxes, per_env_bbox_overrides = self._build_candidate_bboxes(
+            env_bboxes, num_envs, candidates_per_env
+        )
 
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
         for candidate_idx in range(num_candidates):
@@ -315,38 +245,112 @@ def _place_heterogeneous(
             )
 
         all_positions = self._solver.solve(objects, initial_positions, env_bboxes=candidate_bboxes)
+        # solver.solve() populates last_loss_per_env; assert documents the contract.
         assert self._solver.last_loss_per_env is not None
         all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
 
-        ranked_results_per_env: list[list[PlacementResult]] = []
-        for cur_env in range(num_envs):
-            start = cur_env * candidates_per_env
+        candidates: list[PlacementCandidate] = []
+        for candidate_idx in range(num_candidates):
+            cur_env = candidate_idx // candidates_per_env
             env_bbox_overrides = per_env_bbox_overrides[cur_env]
-            env_candidates = [
+            candidates.append(
                 PlacementCandidate(
-                    all_losses[start + idx],
-                    all_positions[start + idx],
-                    self._validate_placement(all_positions[start + idx], bbox_overrides=env_bbox_overrides),
+                    all_losses[candidate_idx],
+                    all_positions[candidate_idx],
+                    self._validate_placement(all_positions[candidate_idx], bbox_overrides=env_bbox_overrides),
                 )
-                for idx in range(candidates_per_env)
-            ]
-            env_candidates.sort(key=lambda candidate: (not candidate.is_valid, candidate.loss))
-            ranked_results_per_env.append([
+            )
+
+        ranked_candidate_slices = self._rank_candidates(candidates, ranking_mode, num_envs, candidates_per_env)
+        ranked_results = [
+            [
                 PlacementResult(
                     success=candidate.is_valid,
                     positions=candidate.positions,
                     final_loss=candidate.loss,
-                    attempts=candidates_per_env,
+                    attempts=attempts_per_result,
                 )
-                for candidate in env_candidates
-            ])
+                for candidate in candidate_slice
+            ]
+            for candidate_slice in ranked_candidate_slices
+        ]
 
-        results = [ranked_results[0] for ranked_results in ranked_results_per_env]
         if self.params.verbose:
-            n_valid = sum(1 for r in results if r.success)
-            print(f"Solved {num_candidates} candidates in one batch (heterogeneous): {n_valid}/{num_envs} env(s) valid")
+            self._print_ranked_summary(ranked_candidate_slices, ranking_mode, num_candidates, num_envs)
 
-        return ranked_results_per_env
+        return ranked_results
+
+    @staticmethod
+    def _build_candidate_bboxes(
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox],
+        num_envs: int,
+        candidates_per_env: int,
+    ) -> tuple[dict[ObjectBase, AxisAlignedBoundingBox], list[dict[ObjectBase, AxisAlignedBoundingBox]]]:
+        """Build solver bboxes with shape (num_envs * candidates_per_env, 3) and per-env views."""
+        per_env_bbox_overrides: list[dict[ObjectBase, AxisAlignedBoundingBox]] = [
+            {
+                obj: AxisAlignedBoundingBox(
+                    min_point=bbox.min_point[cur_env : cur_env + 1],
+                    max_point=bbox.max_point[cur_env : cur_env + 1],
+                )
+                for obj, bbox in env_bboxes.items()
+            }
+            for cur_env in range(num_envs)
+        ]
+
+        candidate_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {}
+        for obj, bbox in env_bboxes.items():
+            candidate_bboxes[obj] = AxisAlignedBoundingBox(
+                min_point=bbox.min_point.repeat_interleave(candidates_per_env, dim=0),
+                max_point=bbox.max_point.repeat_interleave(candidates_per_env, dim=0),
+            )
+
+        return candidate_bboxes, per_env_bbox_overrides
+
+    @staticmethod
+    def _rank_candidates(
+        candidates: list[PlacementCandidate],
+        ranking_mode: Literal["shared_geometry", "env_specific_geometry"],
+        num_envs: int,
+        candidates_per_env: int,
+    ) -> list[list[PlacementCandidate]]:
+        """Return one shared-geometry sorted slice or one sorted slice per env."""
+        if ranking_mode == "shared_geometry":
+            return [sorted(candidates, key=lambda candidate: (not candidate.is_valid, candidate.loss))]
+
+        ranked_candidate_slices: list[list[PlacementCandidate]] = []
+        for cur_env in range(num_envs):
+            start = cur_env * candidates_per_env
+            env_candidates = candidates[start : start + candidates_per_env]
+            ranked_candidate_slices.append(
+                sorted(env_candidates, key=lambda candidate: (not candidate.is_valid, candidate.loss))
+            )
+        return ranked_candidate_slices
+
+    def _print_ranked_summary(
+        self,
+        ranked_candidate_slices: list[list[PlacementCandidate]],
+        ranking_mode: Literal["shared_geometry", "env_specific_geometry"],
+        num_candidates: int,
+        num_envs: int,
+    ) -> None:
+        if ranking_mode == "shared_geometry":
+            candidates = ranked_candidate_slices[0]
+            total_valid = sum(1 for candidate in candidates if candidate.is_valid)
+            finite_losses = [candidate.loss for candidate in candidates if math.isfinite(candidate.loss)]
+            mean_loss = sum(finite_losses) / len(finite_losses) if finite_losses else float("inf")
+            n_valid = sum(1 for candidate in candidates[:num_envs] if candidate.is_valid)
+            print(
+                f"Solved {num_candidates} candidates in one batch: mean loss = {mean_loss:.6f},"
+                f" {total_valid} valid, selected best {num_envs} ({n_valid} valid)"
+            )
+            return
+
+        n_valid = sum(1 for candidate_slice in ranked_candidate_slices if candidate_slice[0].is_valid)
+        print(
+            f"Solved {num_candidates} candidates in one batch (env-specific geometry): "
+            f"{n_valid}/{num_envs} env(s) valid"
+        )
 
     def _generate_initial_positions(
         self,
@@ -365,8 +369,7 @@ def _generate_initial_positions(
             generator: Optional RNG generator for reproducible sampling. When None,
                 uses PyTorch's global RNG.
             child_bboxes: Optional per-object bbox overrides with shape (1, 3).
-                Used by heterogeneous placement to supply the correct variant
-                bbox when computing On-guided initial positions.
+                Used to supply the correct env bbox for On-guided initialization.
 
         Returns:
             Dictionary mapping all objects to their starting positions.
@@ -379,7 +382,12 @@ def _generate_initial_positions(
         positions: dict[ObjectBase, tuple[float, float, float]] = {}
         for obj in objects:
             if obj in anchor_objects:
-                positions[obj] = obj.get_initial_pose().position_xyz
+                initial_pose = obj.get_initial_pose()
+                assert isinstance(initial_pose, Pose), (
+                    f"Anchor object '{obj.name}' must have a fixed Pose before placement, got"
+                    f" {type(initial_pose).__name__}."
+                )
+                positions[obj] = initial_pose.position_xyz
             elif any(isinstance(r, On) for r in obj.get_relations()):
                 bbox_override = child_bboxes.get(obj) if child_bboxes else None
                 positions[obj] = self._compute_on_guided_position(
@@ -400,7 +408,7 @@ def _get_on_parent_world_bbox(
         If the parent is an anchor, return its world bbox directly.
         If the parent is a non-anchor with its own On(anchor) relation, use the anchor's
         world bbox as a proxy. Only one level of indirection is resolved; deeper chains
-        fall back to anchor_bbox. Otherwise fall back to anchor_bbox.
+        fall back to anchor_bbox.
 
         TODO(cvolk): Support full On-relation chains (e.g. spoon -> On(bowl) -> On(plate) -> On(table)).
         """
@@ -450,7 +458,7 @@ def _compute_on_guided_position(
             generator,
         )
 
-        # Z: place child's bottom face at parent top + clearance
+        # Convert from child-origin Z to child-bottom Z so the bottom face lands on the parent top.
         z = float(parent_bbox.max_point[0, 2] + on_relation.clearance_m - child_bbox.min_point[0, 2])
 
         return (x, y, z)
@@ -490,7 +498,7 @@ def _validate_on_relations(
         positions: dict[ObjectBase, tuple[float, float, float]],
         bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> bool:
-        """Validate each On relation; logic matches OnLossStrategy (relation_loss_strategies.py).
+        """Validate each On relation; keep in sync with OnLossStrategy in relation_loss_strategies.py.
 
         1. X: child's footprint entirely within parent's X extent.
         2. Y: child's footprint entirely within parent's Y extent.
@@ -498,8 +506,7 @@ def _validate_on_relations(
 
         Args:
             positions: Solved positions for each object.
-            bbox_overrides: Optional per-object bbox overrides (single-env, shape (1, 3)).
-                Used by heterogeneous placement to supply the correct variant bbox.
+            bbox_overrides: Optional per-object bbox overrides with shape (1, 3).
         """
         for obj in positions:
             for rel in obj.get_relations():
@@ -512,7 +519,6 @@ def _validate_on_relations(
                 parent_bbox = self._resolve_bbox(parent, bbox_overrides)
                 child_world = child_bbox.translated(positions[obj])
                 parent_world = parent_bbox.translated(positions[parent])
-                # 1 & 2: Same as OnLossStrategy X/Y band (child's footprint within parent).
                 if (
                     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]
@@ -522,7 +528,6 @@ def _validate_on_relations(
                     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_bbox.max_point[0, 2].item()
                 child_local_bottom_z: float = child_bbox.min_point[0, 2].item()
                 parent_top_z = parent_local_top_z + positions[parent][2]
@@ -549,31 +554,30 @@ def _validate_no_overlap(
 
         Args:
             positions: Solved positions for each object.
-            bbox_overrides: Optional per-object bbox overrides (single-env, shape (1, 3)).
-                Used by heterogeneous placement to supply the correct variant bbox.
+            bbox_overrides: Optional per-object bbox overrides with shape (1, 3).
         """
-        # Build set of On-related pairs to skip (child, parent) and (parent, child).
         on_pairs: set[tuple] = set()
         anchor_ids: set[int] = set()
         for obj in positions:
             for rel in obj.get_relations():
                 if isinstance(rel, On) and rel.parent in positions:
+                    # The lookup below sees pairs in object-list order, so store
+                    # both directions for symmetric On-pair skipping.
                     on_pairs.add((id(obj), id(rel.parent)))
                     on_pairs.add((id(rel.parent), id(obj)))
             if any(isinstance(r, IsAnchor) for r in obj.get_relations()):
                 anchor_ids.add(id(obj))
 
         clearance_m = self.params.solver_params.clearance_m
+        # Allow tiny residuals from the differentiable solver around the clearance boundary.
         margin = max(0.0, clearance_m - 1e-6)
 
         objects = list(positions.keys())
         for i in range(len(objects)):
             for j in range(i + 1, len(objects)):
                 a, b = objects[i], objects[j]
-                # Skip anchor-anchor pairs (anchors are fixed, solver does not move them).
                 if id(a) in anchor_ids and id(b) in anchor_ids:
                     continue
-                # Pairs related by an On relation are excluded from the overlap check.
                 if (id(a), id(b)) in on_pairs:
                     continue
 
@@ -597,8 +601,7 @@ def _validate_placement(
 
         Args:
             positions: Dictionary mapping objects to their solved (x, y, z) positions.
-            bbox_overrides: Optional per-object bbox overrides (single-env, shape (1, 3)).
-                Used by heterogeneous placement to supply the correct variant bbox.
+            bbox_overrides: Optional per-object bbox overrides with shape (1, 3).
 
         Returns:
             True if no overlaps exist and On relations hold, False otherwise.
@@ -621,9 +624,7 @@ def _apply_positions(
         Rotation is taken from RotateAroundSolution marker if present, otherwise identity.
         """
         num_envs = len(positions_per_env)
-        # Objects are the same for every environment. Extract them.
         objects = list(positions_per_env[0])
-        # Apply pose for each object.
         for obj in objects:
             if obj in anchor_objects:
                 continue
@@ -646,28 +647,12 @@ def _apply_positions(
                 obj.set_initial_pose(PosePerEnv(poses=poses))
 
     def _get_random_around_solution(self, obj: ObjectBase) -> RandomAroundSolution | None:
-        """Get RandomAroundSolution marker from object if present.
-
-        Args:
-            obj: Object to check for the marker.
-
-        Returns:
-            The RandomAroundSolution marker if found, None otherwise.
-        """
         for rel in obj.get_relations():
             if isinstance(rel, RandomAroundSolution):
                 return rel
         return None
 
     def _get_rotate_around_solution(self, obj: ObjectBase) -> RotateAroundSolution | None:
-        """Get RotateAroundSolution marker from object if present.
-
-        Args:
-            obj: Object to check for the marker.
-
-        Returns:
-            The RotateAroundSolution marker if found, None otherwise.
-        """
         for rel in obj.get_relations():
             if isinstance(rel, RotateAroundSolution):
                 return rel
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index a43d69e81..3ee60bb31 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -7,6 +7,7 @@
 
 import random
 import torch
+from dataclasses import replace
 from typing import TYPE_CHECKING
 
 from isaaclab_arena.relations.bounding_box_helpers import has_heterogeneous_objects
@@ -59,7 +60,9 @@ def __init__(
 
         # 2. Configure dependencies and per-env storage.
         self._objects = list(objects)
-        self._placer = ObjectPlacer(params=placer_params)
+        # Pool construction ranks several candidate layouts per env and applies
+        # poses only when a sampled layout is used.
+        self._placer = ObjectPlacer(params=replace(placer_params, apply_positions_to_objects=False))
         self._pool_size = pool_size
         self._had_fallbacks = False
         self._layout_pools: dict[int, list[PlacementResult]] = {cur_env: [] for cur_env in range(self._num_envs)}
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index e0f786cbb..032387c7c 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -358,6 +358,16 @@ def test_object_placer_place_ranked_per_env_returns_sorted_env_lists():
     with pytest.raises(AssertionError):
         placer.place_ranked_per_env([desk, hetero], num_envs=3, results_per_env=0)
 
+    apply_params = ObjectPlacerParams(
+        solver_params=solver_params,
+        apply_positions_to_objects=True,
+    )
+    apply_results = ObjectPlacer(params=apply_params).place_ranked_per_env(
+        [desk, hetero], num_envs=3, results_per_env=1
+    )
+    assert len(apply_results) == 3
+    assert hetero.get_initial_pose() is None
+
 
 def test_object_placer_homogeneous_path_returns_multi_env_result():
     """When no heterogeneous objects exist, the homogeneous path is used."""

From 24acf31a23b39d5f3bfe41b2a5bf2a0ac6524586 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 26 May 2026 11:53:32 -0700
Subject: [PATCH 39/46] refactor and address comments

---
 isaaclab_arena/assets/object_set.py           |  53 +++++---
 .../environments/relation_solver_interface.py |  43 +++---
 .../relations/bounding_box_helpers.py         |  23 +++-
 isaaclab_arena/relations/object_placer.py     |  61 ++++++---
 isaaclab_arena/relations/placement_events.py  |  20 ++-
 .../relations/pooled_object_placer.py         | 122 ++++++++++++------
 .../tests/test_heterogeneous_placement.py     | 100 +++++++++++---
 .../tests/test_object_placer_init.py          |  27 ++++
 isaaclab_arena/tests/test_object_set.py       |  38 +++++-
 isaaclab_arena/tests/test_placement_events.py |  47 ++++++-
 10 files changed, 391 insertions(+), 143 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index d183dfd87..5f1c19598 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -19,9 +19,7 @@
 
 
 class RigidObjectSet(Object):
-    """
-    A set of rigid objects.
-    """
+    """A set of rigid objects with one member selected per environment."""
 
     def __init__(
         self,
@@ -97,26 +95,27 @@ def __init__(
 
     @property
     def object_usd_paths(self) -> list[str]:
-        """Per-env USD paths passed to MultiUsdFileCfg.
+        """USD paths passed to MultiUsdFileCfg.
 
-        Before variant indices are assigned, returns the member list.
-        After assignment, returns one path per env based on variant_indices_by_env.
+        Before assignment this is the member USD list. After assignment this
+        returns one USD path per environment based on ``variant_indices_by_env``.
         """
         if self.variant_indices_by_env is not None:
             return [self.member_usd_paths[idx] for idx in self.variant_indices_by_env]
         return self.member_usd_paths
 
     def get_bounding_box(self) -> AxisAlignedBoundingBox:
-        """Get the bounding box of the object set.
+        """Return one conservative local bbox for callers that cannot vary by env.
 
-        This compatibility fallback returns the member bbox with the greatest
-        z-extent. Heterogeneous placement should use get_bounding_box_per_env()
-        after assign_variants() so each env uses its actual variant geometry.
+        The returned bbox has shape ``(1, 3)`` and uses the member with the
+        greatest z-extent. Heterogeneous placement uses
+        ``get_bounding_box_per_env()`` after ``assign_variants()`` so each env
+        uses its actual variant geometry.
         """
         return max(self.objects, key=lambda obj: obj.get_bounding_box().size[0, 2].item()).get_bounding_box()
 
-    def assign_variants(self, num_envs: int) -> None:
-        """Assign one member-variant index per environment.
+    def assign_variants(self, num_envs: int, variant_seed: int | None = None) -> None:
+        """Fix one member-variant index per environment.
 
         The assignment is fixed for the lifetime of the object set: subsequent
         calls with the same ``num_envs`` are no-ops, and a call with a
@@ -124,12 +123,12 @@ def assign_variants(self, num_envs: int) -> None:
         env independently samples one variant; otherwise assignments repeat
         the member order across environments.
 
-        Callers (typically the placer once ``num_envs`` is known) must invoke
-        this before reading ``variant_indices_by_env`` or
-        ``get_bounding_box_per_env``.
+        Callers invoke this once ``num_envs`` is known, before reading
+        ``variant_indices_by_env`` or ``get_bounding_box_per_env``.
 
         Args:
             num_envs: Number of environments to assign variants for.
+            variant_seed: Optional seed used when random_choice=True.
         """
         if self.variant_indices_by_env is not None:
             assert len(self.variant_indices_by_env) == num_envs, (
@@ -137,16 +136,17 @@ def assign_variants(self, num_envs: int) -> None:
                 f"{len(self.variant_indices_by_env)} envs, got request for {num_envs}."
             )
             return
-        self._set_variant_indices_by_env(self._generate_variant_indices(num_envs))
+        self._set_variant_indices_by_env(self._generate_variant_indices(num_envs, variant_seed=variant_seed))
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
-        """Get the actual bounding box for each env's variant.
+        """Return the local bbox for each env's assigned variant.
 
         Unlike the single-bbox compatibility fallback, this returns the real
         local bbox of the variant assigned to each env, enabling correct
         collision-free placement for heterogeneous scenes.
 
-        Requires ``assign_variants(num_envs)`` to have been called first.
+        Requires ``assign_variants(num_envs)`` to have been called first. The
+        returned bbox has shape ``(num_envs, 3)``.
 
         Args:
             num_envs: Number of environments. Must match the assignment.
@@ -174,19 +174,30 @@ def get_contact_sensor_cfg(self, contact_against_object: ObjectBase | None = Non
         # and we can use the canonical first member USD to find the shallowest rigid body.
         return super().get_contact_sensor_cfg(contact_against_object, usd_path=self.member_usd_paths[0])
 
-    def _generate_variant_indices(self, num_envs: int) -> list[int]:
+    def _generate_variant_indices(self, num_envs: int, variant_seed: int | None = None) -> list[int]:
+        """Return one member index per env.
+
+        Ordered sets repeat member order. Random sets sample independently per
+        env, using a local generator when variant_seed is set.
+        """
         n = len(self.objects)
         if not self.random_choice:
             return [env_idx % n for env_idx in range(num_envs)]
-        return torch.randint(low=0, high=n, size=(num_envs,)).tolist()
+        if variant_seed is None:
+            return torch.randint(low=0, high=n, size=(num_envs,)).tolist()
+        generator = torch.Generator()
+        generator.manual_seed(variant_seed)
+        return torch.randint(low=0, high=n, size=(num_envs,), generator=generator).tolist()
 
     def _set_variant_indices_by_env(self, variant_indices_by_env: list[int]) -> None:
-        """Validate and store variant indices, then sync the spawn config's USD path list."""
+        """Validate and store variant indices, then sync spawn config when it exists."""
         n = len(self.objects)
         assert all(
             0 <= idx < n for idx in variant_indices_by_env
         ), f"RigidObjectSet '{self.name}' variant indices must be in [0, {n}); got {variant_indices_by_env}."
         self.variant_indices_by_env = list(variant_indices_by_env)
+        # During __init__, Object.object_cfg has not been built yet; _generate_rigid_cfg()
+        # reads object_usd_paths after this assignment.
         spawn_cfg = self.object_cfg.spawn if getattr(self, "object_cfg", None) is not None else None
         if isinstance(spawn_cfg, sim_utils.MultiUsdFileCfg):
             spawn_cfg.usd_path = self.object_usd_paths
diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py
index 8f6bf6bf5..5742fa6a5 100644
--- a/isaaclab_arena/environments/relation_solver_interface.py
+++ b/isaaclab_arena/environments/relation_solver_interface.py
@@ -53,14 +53,19 @@ def solve_and_apply_relation_placement(
     if resolve_on_reset is not None:
         placer_params.resolve_on_reset = resolve_on_reset
 
-    # TODO(xinjieyao, 2026-05-22): Add joint object/embodiment placement once task-dependent
-    # reachability constraints are available. For now this always uses the object-only placer.
     placement_pool = PooledObjectPlacer(
         objects=objects,
         placer_params=placer_params,
         pool_size=num_envs * placer_params.min_unique_layouts_per_env,
         num_envs=num_envs,
     )
+
+    if placement_pool.had_fallbacks:
+        print(
+            "Warning: Relation placement pool accepted best-loss fallback layouts "
+            "that failed strict placement validation."
+        )
+
     return _apply_relation_placement_result(
         objects=objects,
         placer_params=placer_params,
@@ -77,10 +82,8 @@ def _apply_relation_placement_result(
 ) -> EventTermCfg | None:
     """Apply selected layouts to object spawn state and build reset event config."""
     anchor_objects_set = set(get_anchor_objects(objects))
-    # Prevent external pose-reset events from conflicting with relation-solved objects.
     _validate_no_conflicting_pose_reset_events(objects, anchor_objects_set)
 
-    # Anchor objects do not move, so no need to apply reset event.
     if anchor_objects_set == set(objects):
         return None
 
@@ -108,18 +111,21 @@ def _apply_dynamic_spawn_pose(
     """Set initial spawn pose from one layout and return the reset placement event."""
     from isaaclab.managers import EventTermCfg
 
-    layout = placement_pool.sample_with_replacement(1)[0]
-    for obj in objects:
-        if obj in anchor_objects_set:
-            continue
-        pos = layout.positions.get(obj)
-        if pos is None:
-            continue
-        object_cfg = getattr(obj, "object_cfg", None)
-        if object_cfg is None:
-            raise RuntimeError(f"Object '{obj.name}' must have object_cfg initialized before placement.")
-        object_cfg.init_state.pos = pos
-        object_cfg.init_state.rot = get_rotation_xyzw(obj)
+    if placement_pool.requires_env_indexed_layouts:
+        print("Warning: Skipping static init_state seeding for env-indexed placement layouts.")
+    else:
+        layout = placement_pool.sample_with_replacement(1)[0]
+        for obj in objects:
+            if obj in anchor_objects_set:
+                continue
+            pos = layout.positions.get(obj)
+            if pos is None:
+                raise RuntimeError(f"Pool layout is missing object '{obj.name}'.")
+            object_cfg = getattr(obj, "object_cfg", None)
+            if object_cfg is None:
+                raise RuntimeError(f"Object '{obj.name}' must have object_cfg initialized before placement.")
+            object_cfg.init_state.pos = pos
+            object_cfg.init_state.rot = get_rotation_xyzw(obj)
 
     return EventTermCfg(
         func=solve_and_place_objects,
@@ -147,10 +153,9 @@ def _apply_static_initial_poses(
         for env_idx in range(num_envs):
             pos = layouts[env_idx].positions.get(obj)
             if pos is None:
-                break
+                raise RuntimeError(f"Placement layout for env {env_idx} is missing object '{obj.name}'.")
             poses.append(Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw))
-        else:
-            obj.set_initial_pose(PosePerEnv(poses=poses))
+        obj.set_initial_pose(PosePerEnv(poses=poses))
 
 
 def _validate_no_conflicting_pose_reset_events(
diff --git a/isaaclab_arena/relations/bounding_box_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py
index 82093ca14..6f466b556 100644
--- a/isaaclab_arena/relations/bounding_box_helpers.py
+++ b/isaaclab_arena/relations/bounding_box_helpers.py
@@ -17,23 +17,30 @@
 if TYPE_CHECKING:
     from isaaclab_arena.assets.object_base import ObjectBase
 
+VARIANT_SEED_STRIDE = 1_000_003
+
 
 def has_heterogeneous_objects(objects: list[ObjectBase]) -> bool:
-    """True if any object in the list is a RigidObjectSet."""
+    """Return whether placement must use env-specific object geometry."""
     from isaaclab_arena.assets.object_set import RigidObjectSet
 
     return any(isinstance(obj, RigidObjectSet) for obj in objects)
 
 
-def assign_variants_for_envs(objects: list[ObjectBase], num_envs: int) -> None:
+def assign_variants_for_envs(objects: list[ObjectBase], num_envs: int, placement_seed: int | None = None) -> None:
     """Assign per-env variants on every RigidObjectSet in the list.
 
-    Placers call this at the boundary before any per-env geometry reads.
-    Objects without variants are ignored.
+    Placers call this once they know the real environment count, before
+    requesting per-env bounding boxes. Objects without variants are ignored.
     """
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+
+    variant_set_idx = 0
     for obj in objects:
-        if hasattr(obj, "assign_variants"):
-            obj.assign_variants(num_envs)
+        if isinstance(obj, RigidObjectSet):
+            variant_seed = None if placement_seed is None else placement_seed + VARIANT_SEED_STRIDE * variant_set_idx
+            obj.assign_variants(num_envs, variant_seed=variant_seed)
+            variant_set_idx += 1
 
 
 def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBoundingBox:
@@ -42,7 +49,9 @@ def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBound
     RigidObjectSet delegates to its own get_bounding_box_per_env.
     All other objects broadcast their single bbox.
     """
-    if hasattr(obj, "get_bounding_box_per_env"):
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+
+    if isinstance(obj, RigidObjectSet):
         return obj.get_bounding_box_per_env(num_envs)
 
     bbox = obj.get_bounding_box()
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index dfdb3ccc9..6319bc1d4 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -94,22 +94,20 @@ def place(
                 placement (one layout per env).
             result_per_env: When True (default), each environment gets a distinct
                 layout. When False, homogeneous objects use a single best layout
-                for all environments. Ignored for heterogeneous scenes; those
-                always produce one layout per env so each env is solved with its
-                assigned variant geometry.
+                for all environments. Heterogeneous object sets always return
+                one layout per env, even when this flag is False, because each
+                env must be solved against its assigned variant geometry.
 
         Returns:
-            PlacementResult when a single layout is produced; MultiEnvPlacementResult
-            when multiple per-env layouts are produced.
+            ``PlacementResult`` for one produced layout, otherwise
+            ``MultiEnvPlacementResult``. A heterogeneous multi-env call produces
+            ``MultiEnvPlacementResult`` even when ``result_per_env`` is False.
         """
         anchor_objects_set, generator = self._prepare_placement(objects)
 
         uses_env_specific_bboxes = has_heterogeneous_objects(objects)
-        # Env-specific variants cannot share one solved layout because each env
-        # must be solved against its assigned object geometry.
         num_results = num_envs if result_per_env or uses_env_specific_bboxes else 1
         max_attempts = self.params.max_placement_attempts
-        # Solve max_attempts candidates per result so we can pick the best survivor.
         if uses_env_specific_bboxes:
             ranked_results = self._place_ranked(
                 objects,
@@ -137,7 +135,6 @@ def place(
         if self.params.apply_positions_to_objects:
             self._apply_positions(positions_per_env, anchor_objects_set)
 
-        # Single layout when callers don't need per-env results; otherwise return per-env layouts.
         if num_results == 1:
             return results_per_env[0]
         return MultiEnvPlacementResult(results=results_per_env)
@@ -160,7 +157,6 @@ def place_ranked_per_env(
         assert results_per_env > 0, f"results_per_env must be positive, got {results_per_env}"
         anchor_objects_set, generator = self._prepare_placement(objects)
         max_attempts = self.params.max_placement_attempts
-        # Callers ask for independent candidate lists per env, not one shared-geometry top-N list.
         ranked_results_per_env = self._place_ranked(
             objects,
             anchor_objects_set,
@@ -226,7 +222,7 @@ def _place_ranked(
                 " variants."
             )
 
-        assign_variants_for_envs(objects, num_envs)
+        assign_variants_for_envs(objects, num_envs, placement_seed=self.params.placement_seed)
         env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
         num_candidates = num_envs * candidates_per_env
         candidate_bboxes, per_env_bbox_overrides = self._build_candidate_bboxes(
@@ -245,7 +241,6 @@ def _place_ranked(
             )
 
         all_positions = self._solver.solve(objects, initial_positions, env_bboxes=candidate_bboxes)
-        # solver.solve() populates last_loss_per_env; assert documents the contract.
         assert self._solver.last_loss_per_env is not None
         all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
 
@@ -375,7 +370,7 @@ def _generate_initial_positions(
             Dictionary mapping all objects to their starting positions.
         """
         first_anchor = next(obj for obj in objects if obj in anchor_objects)
-        anchor_bbox = first_anchor.get_world_bounding_box()
+        anchor_bbox = self._get_world_bbox_for_init(first_anchor, child_bboxes)
 
         cx, cy, cz = float(anchor_bbox.center[0, 0]), float(anchor_bbox.center[0, 1]), float(anchor_bbox.center[0, 2])
 
@@ -383,25 +378,44 @@ def _generate_initial_positions(
         for obj in objects:
             if obj in anchor_objects:
                 initial_pose = obj.get_initial_pose()
-                assert isinstance(initial_pose, Pose), (
-                    f"Anchor object '{obj.name}' must have a fixed Pose before placement, got"
-                    f" {type(initial_pose).__name__}."
-                )
+                if not isinstance(initial_pose, Pose):
+                    raise TypeError(
+                        f"Anchor object '{obj.name}' must have a fixed Pose before placement, got"
+                        f" {type(initial_pose).__name__}."
+                    )
                 positions[obj] = initial_pose.position_xyz
             elif any(isinstance(r, On) for r in obj.get_relations()):
                 bbox_override = child_bboxes.get(obj) if child_bboxes else None
                 positions[obj] = self._compute_on_guided_position(
-                    obj, anchor_objects, anchor_bbox, generator, child_bbox=bbox_override
+                    obj, anchor_objects, anchor_bbox, generator, child_bbox=bbox_override, bbox_overrides=child_bboxes
                 )
             else:
                 positions[obj] = (cx, cy, cz)
         return positions
 
+    @staticmethod
+    def _get_world_bbox_for_init(
+        obj: ObjectBase,
+        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None,
+    ) -> AxisAlignedBoundingBox:
+        bbox_override = bbox_overrides.get(obj) if bbox_overrides else None
+        if bbox_override is None:
+            return obj.get_world_bounding_box()
+
+        initial_pose = obj.get_initial_pose()
+        if not isinstance(initial_pose, Pose):
+            raise TypeError(
+                f"Object '{obj.name}' must have a fixed Pose to use an env-specific world bbox,"
+                f" got {type(initial_pose).__name__}."
+            )
+        return bbox_override.translated(initial_pose.position_xyz)
+
     def _get_on_parent_world_bbox(
         self,
         parent: ObjectBase,
         anchor_objects: set[ObjectBase],
         anchor_bbox: AxisAlignedBoundingBox,
+        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> AxisAlignedBoundingBox:
         """Resolve the world bbox of an On relation's parent for initialization purposes.
 
@@ -413,10 +427,10 @@ def _get_on_parent_world_bbox(
         TODO(cvolk): Support full On-relation chains (e.g. spoon -> On(bowl) -> On(plate) -> On(table)).
         """
         if parent in anchor_objects:
-            return parent.get_world_bounding_box()
+            return self._get_world_bbox_for_init(parent, bbox_overrides)
         for rel in parent.get_relations():
             if isinstance(rel, On) and rel.parent in anchor_objects:
-                return rel.parent.get_world_bounding_box()
+                return self._get_world_bbox_for_init(rel.parent, bbox_overrides)
         return anchor_bbox
 
     def _compute_on_guided_position(
@@ -426,6 +440,7 @@ def _compute_on_guided_position(
         anchor_bbox: AxisAlignedBoundingBox,
         generator: torch.Generator | None = None,
         child_bbox: AxisAlignedBoundingBox | None = None,
+        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> tuple[float, float, float]:
         """Compute an initial position for an object with an On relation.
 
@@ -437,9 +452,13 @@ def _compute_on_guided_position(
                 uses PyTorch's global RNG.
             child_bbox: Optional bbox override for the child object. When None,
                 obj.get_bounding_box() is used.
+            bbox_overrides: Optional per-object bbox overrides for resolving
+                heterogeneous On parents.
         """
         on_relation = next(r for r in obj.get_relations() if isinstance(r, On))
-        parent_bbox = self._get_on_parent_world_bbox(on_relation.parent, anchor_objects, anchor_bbox)
+        parent_bbox = self._get_on_parent_world_bbox(
+            on_relation.parent, anchor_objects, anchor_bbox, bbox_overrides=bbox_overrides
+        )
         if child_bbox is None:
             child_bbox = obj.get_bounding_box()
 
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index 5d94ad72a..86056ab34 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -34,10 +34,10 @@ def solve_and_place_objects(
 ) -> None:
     """Coordinated reset event that draws layouts from the pool and writes poses.
 
-    Registered as a single EventTermCfg(mode="reset"). Env-specific
-    layouts advance by one full env round so each result still matches its
-    absolute env id. Reusable layouts draw only for the environments being
-    reset.
+    Registered as a single EventTermCfg(mode="reset"). Env-indexed pools
+    consume one complete scene round, then write only the resetting subset so
+    each result still matches its absolute env id. Reusable pools consume only
+    the number of resetting envs because those layouts are interchangeable.
 
     Args:
         env: The Isaac Lab environment.
@@ -51,9 +51,10 @@ def solve_and_place_objects(
     reset_env_ids = env_ids.tolist()
     if placement_pool.requires_env_indexed_layouts:
         num_scene_envs = env.scene.env_origins.shape[0]
-        assert (
-            placement_pool.num_envs == num_scene_envs
-        ), f"Placement pool has {placement_pool.num_envs} envs, but scene has {num_scene_envs} env origins."
+        if placement_pool.num_envs != num_scene_envs:
+            raise ValueError(
+                f"Placement pool has {placement_pool.num_envs} envs, but scene has {num_scene_envs} env origins."
+            )
         all_results = placement_pool.sample_without_replacement(num_scene_envs)
         results_by_env = {cur_env: all_results[cur_env] for cur_env in reset_env_ids}
     else:
@@ -67,6 +68,11 @@ def solve_and_place_objects(
     for cur_env in reset_env_ids:
         env_id_tensor = torch.tensor([cur_env], device=env.device)
         result = results_by_env[cur_env]
+        if not result.success:
+            print(
+                "Warning: Writing best-loss fallback placement for "
+                f"env {cur_env}; layout failed strict placement validation."
+            )
         positions = result.positions
         for obj, pos in positions.items():
             if obj in anchor_objects_set:
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 3ee60bb31..a168ea611 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -7,7 +7,7 @@
 
 import random
 import torch
-from dataclasses import replace
+from dataclasses import dataclass, replace
 from typing import TYPE_CHECKING
 
 from isaaclab_arena.relations.bounding_box_helpers import has_heterogeneous_objects
@@ -19,20 +19,54 @@
     from isaaclab_arena.assets.object_base import ObjectBase
 
 
+@dataclass
+class EnvLayoutPool:
+    """Unread layout queue for one absolute environment."""
+
+    layouts: list[PlacementResult]
+    cursor: int = 0
+
+    @property
+    def available(self) -> int:
+        return len(self.layouts) - self.cursor
+
+    def discard_consumed(self) -> None:
+        self.layouts = self.layouts[self.cursor :]
+        self.cursor = 0
+
+    def append(self, layout: PlacementResult) -> None:
+        self.layouts.append(layout)
+
+    def extend(self, layouts: list[PlacementResult]) -> None:
+        self.layouts.extend(layouts)
+
+    def next(self) -> PlacementResult:
+        if self.cursor >= len(self.layouts):
+            raise IndexError("No unread layouts remain in this env pool.")
+        layout = self.layouts[self.cursor]
+        self.cursor += 1
+        return layout
+
+
 class PooledObjectPlacer:
-    """Object placer that keeps a pool of optimized layouts.
+    """Object placer that maintains solved placement layouts.
+
+    Storage is organized as one queue per environment. Env-specific layouts
+    are solved for a fixed env's object geometry and must be consumed in
+    complete env rounds. Reusable layouts are interchangeable and can be
+    consumed one at a time from the pooled queues.
 
-    Storage: num_envs independent layout pools, each with its own read
-    cursor. Env-specific layouts are solved
-    against a fixed env's object geometry and must be sampled in complete env
-    rounds. Reusable layouts can be consumed one at a time.
+    Strictly valid layouts are preferred. If no valid layout is available for
+    a solve batch, the best-loss solver result is kept as a visible fallback.
 
     The pool is refilled automatically when an env's queue runs out.
 
-    * sample_without_replacement — returns the next count layouts.
-      Env-specific layouts require count to be a multiple of num_envs.
-    * sample_with_replacement — picks count layouts at random per
-      env-slot (non-consuming). Used for static initial positions.
+    * ``sample_without_replacement(count)`` consumes ``count`` layouts. For
+      env-specific layouts, ``count`` must be a multiple of ``num_envs`` and
+      may cover multiple complete env rounds.
+    * ``sample_with_replacement(count)`` is non-consuming. Env-specific layouts
+      are sampled from matching env slots; reusable layouts are sampled IID from
+      all stored layouts.
 
     Args:
         objects: All objects (including anchors) participating in relation solving.
@@ -52,25 +86,25 @@ def __init__(
         if pool_size < 1:
             raise ValueError(f"pool_size must be >= 1, got {pool_size}")
         self._uses_env_specific_bboxes = has_heterogeneous_objects(objects)
-        if self._uses_env_specific_bboxes:
-            assert num_envs is not None, "num_envs is required when layouts use env-specific object variants."
+        if self._uses_env_specific_bboxes and num_envs is None:
+            raise ValueError("num_envs is required when layouts use env-specific object variants.")
         self._num_envs = num_envs if num_envs is not None else 1
         if self._num_envs < 1:
             raise ValueError(f"num_envs must be >= 1, got {self._num_envs}")
 
-        # 2. Configure dependencies and per-env storage.
         self._objects = list(objects)
         # Pool construction ranks several candidate layouts per env and applies
         # poses only when a sampled layout is used.
         self._placer = ObjectPlacer(params=replace(placer_params, apply_positions_to_objects=False))
         self._pool_size = pool_size
         self._had_fallbacks = False
-        self._layout_pools: dict[int, list[PlacementResult]] = {cur_env: [] for cur_env in range(self._num_envs)}
-        self._layout_cursors: dict[int, int] = {cur_env: 0 for cur_env in range(self._num_envs)}
+        self._base_placement_seed = placer_params.placement_seed
+        self._next_seed_offset = 0
+        self._env_pools: list[EnvLayoutPool] = [EnvLayoutPool([]) for _ in range(self._num_envs)]
 
         self._solve_and_store(pool_size)
-        for cur_env, pool in self._layout_pools.items():
-            if not pool:
+        for cur_env, pool in enumerate(self._env_pools):
+            if not pool.layouts:
                 raise RuntimeError(
                     f"Placement pool failed to produce any valid layouts for env {cur_env} "
                     f"from {pool_size} attempts. Check object relations and constraints."
@@ -82,7 +116,7 @@ def __init__(
 
     def _available_per_env(self) -> list[int]:
         """Number of unread layouts in each env's pool (length num_envs)."""
-        return [len(self._layout_pools[cur_env]) - self._layout_cursors[cur_env] for cur_env in range(self._num_envs)]
+        return [pool.available for pool in self._env_pools]
 
     def _total_available(self) -> int:
         """Total unread layouts across all env pools."""
@@ -90,10 +124,15 @@ def _total_available(self) -> int:
 
     def _discard_consumed_layouts(self) -> None:
         """Drop consumed layouts from every env pool before appending new layouts."""
-        for cur_env in self._layout_pools:
-            idx = self._layout_cursors[cur_env]
-            self._layout_pools[cur_env] = self._layout_pools[cur_env][idx:]
-            self._layout_cursors[cur_env] = 0
+        for pool in self._env_pools:
+            pool.discard_consumed()
+
+    def _prepare_seeded_solve(self, num_candidates: int) -> None:
+        """Avoid replaying the same candidate sequence on seeded refills."""
+        if self._base_placement_seed is None:
+            return
+        self._placer.params.placement_seed = self._base_placement_seed + self._next_seed_offset
+        self._next_seed_offset += num_candidates
 
     def _solve_and_store(self, num_layouts: int) -> None:
         """Solve layouts in batches until every env has target_per_env unread layouts.
@@ -134,6 +173,7 @@ def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
         If no candidate passes strict validation, fall back to best-loss results
         to match the pre-pool behavior used by existing environments.
         """
+        self._prepare_seeded_solve(num_layouts * self._placer.params.max_placement_attempts)
         with torch.inference_mode(False):
             result = self._placer.place(self._objects, num_envs=num_layouts, result_per_env=True)
 
@@ -168,7 +208,7 @@ def _store_reusable_results(self, layouts: list[PlacementResult]) -> None:
         available = self._available_per_env()
         for layout in layouts:
             cur_env = min(range(self._num_envs), key=available.__getitem__)
-            self._layout_pools[cur_env].append(layout)
+            self._env_pools[cur_env].append(layout)
             available[cur_env] += 1
 
     def _solve_env_ranked_layouts(self, num_layouts: int) -> tuple[list[list[PlacementResult]], int]:
@@ -179,6 +219,10 @@ def _solve_env_ranked_layouts(self, num_layouts: int) -> tuple[list[list[Placeme
         are extra environments.
         """
         layouts_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
+        # ObjectPlacer seeds each candidate as placement_seed + candidate_idx.
+        # Keep this count aligned with place_ranked_per_env's candidate layout.
+        num_candidates = self._num_envs * layouts_per_env * self._placer.params.max_placement_attempts
+        self._prepare_seeded_solve(num_candidates)
 
         with torch.inference_mode(False):
             ranked_results_per_env = self._placer.place_ranked_per_env(
@@ -204,17 +248,17 @@ def _store_env_matched_results(
             env_results = ranked_results_per_env[cur_env][:layouts_per_env]
             valid_results = [r for r in env_results if r.success]
             if valid_results:
-                self._layout_pools[cur_env].extend(valid_results)
+                self._env_pools[cur_env].extend(valid_results)
                 total_valid += len(valid_results)
             else:
-                self._layout_pools[cur_env].extend(env_results)
+                self._env_pools[cur_env].extend(env_results)
                 fallback_envs.append(cur_env)
                 self._had_fallbacks = True
 
         total_solved = sum(min(len(env_results), layouts_per_env) for env_results in ranked_results_per_env)
         if total_valid < total_solved:
             msg = (
-                f"Placement pool (env-specific bbox layouts): solved {total_solved} candidates,"
+                f"Warning: Placement pool (env-specific bbox layouts) solved {total_solved} candidates,"
                 f" {total_valid} valid, {total_solved - total_valid} failed validation"
             )
             if fallback_envs:
@@ -226,12 +270,12 @@ def _store_env_matched_results(
     # ------------------------------------------------------------------
 
     def sample_without_replacement(self, count: int) -> list[PlacementResult]:
-        """Return the next count layouts.
+        """Return the next ``count`` layouts.
 
         Env-specific layouts are returned as complete rounds of
-        [env_0, env_1, ..., env_{num_envs-1}] so each result still maps
-        to the absolute environment it was solved for. Reusable layouts are
-        interchangeable and consume only count entries.
+        ``[env_0, env_1, ..., env_{num_envs-1}]`` so each result still maps
+        to the absolute environment it was solved for. Reusable layouts consume
+        exactly ``count`` interchangeable entries.
 
         Args:
             count: Number of layouts to return.
@@ -256,14 +300,13 @@ def _sample_env_indexed_without_replacement(self, count: int) -> list[PlacementR
         results: list[PlacementResult] = []
         for _ in range(layouts_per_env):
             for cur_env in range(self._num_envs):
-                idx = self._layout_cursors[cur_env]
-                if idx >= len(self._layout_pools[cur_env]):
+                pool = self._env_pools[cur_env]
+                if pool.available <= 0:
                     raise RuntimeError(
                         f"Placement pool: env {cur_env} has no more valid layouts. "
                         "The solver is not producing enough valid placements."
                     )
-                results.append(self._layout_pools[cur_env][idx])
-                self._layout_cursors[cur_env] = idx + 1
+                results.append(pool.next())
         return results
 
     def _sample_reusable_without_replacement(self, count: int) -> list[PlacementResult]:
@@ -281,14 +324,13 @@ def _sample_reusable_without_replacement(self, count: int) -> list[PlacementResu
         results: list[PlacementResult] = []
         for _ in range(count):
             cur_env = max(range(self._num_envs), key=available.__getitem__)
-            idx = self._layout_cursors[cur_env]
-            if idx >= len(self._layout_pools[cur_env]):
+            pool = self._env_pools[cur_env]
+            if pool.available <= 0:
                 raise RuntimeError(
                     f"Placement pool: env {cur_env} has no more valid layouts. "
                     "The solver is not producing enough valid placements."
                 )
-            results.append(self._layout_pools[cur_env][idx])
-            self._layout_cursors[cur_env] = idx + 1
+            results.append(pool.next())
             available[cur_env] -= 1
         return results
 
@@ -318,11 +360,11 @@ def sample_with_replacement(self, count: int) -> list[PlacementResult]:
             results: list[PlacementResult] = []
             for i in range(count):
                 cur_env = i % self._num_envs
-                pool = self._layout_pools[cur_env]
+                pool = self._env_pools[cur_env].layouts
                 assert pool, f"Env {cur_env} has no valid layouts to sample from."
                 results.append(random.choice(pool))
             return results
-        all_layouts = [layout for pool in self._layout_pools.values() for layout in pool]
+        all_layouts = [layout for pool in self._env_pools for layout in pool.layouts]
         return random.choices(all_layouts, k=count)
 
     @property
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 032387c7c..a35ab9b7e 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -32,24 +32,38 @@
 
 @pytest.fixture(autouse=True)
 def _patch_bounding_box_helpers_for_test_doubles(monkeypatch):
-    """Allow HeterogeneousDummyObject to trigger the heterogeneous placement path.
+    """Let ``HeterogeneousDummyObject`` flow through the heterogeneous placement path.
 
-    Only has_heterogeneous_objects needs patching — it uses isinstance(obj, RigidObjectSet)
-    which won't match test doubles. The downstream functions (assign_variants_for_envs,
-    get_bounding_box_per_env) already duck-type via hasattr.
+    Production dispatch uses ``isinstance(RigidObjectSet)``, but these tests use
+    lightweight ``DummyObject`` subclasses with ``get_bounding_box_per_env(...)``.
+    Patch every module that binds the helper by name at import time.
     """
     from isaaclab_arena.relations import bounding_box_helpers
 
-    _original_has_het = bounding_box_helpers.has_heterogeneous_objects
+    original_has_het = bounding_box_helpers.has_heterogeneous_objects
+    original_bbox_per_env = bounding_box_helpers.get_bounding_box_per_env
 
-    def _patched_has_het(objects):
-        if _original_has_het(objects):
-            return True
-        return any(hasattr(obj, "get_bounding_box_per_env") for obj in objects)
+    def has_het_with_doubles(objects):
+        return original_has_het(objects) or any(hasattr(obj, "get_bounding_box_per_env") for obj in objects)
 
-    monkeypatch.setattr("isaaclab_arena.relations.bounding_box_helpers.has_heterogeneous_objects", _patched_has_het)
-    monkeypatch.setattr("isaaclab_arena.relations.object_placer.has_heterogeneous_objects", _patched_has_het)
-    monkeypatch.setattr("isaaclab_arena.relations.pooled_object_placer.has_heterogeneous_objects", _patched_has_het)
+    def bbox_per_env_with_doubles(obj, num_envs):
+        if hasattr(obj, "get_bounding_box_per_env"):
+            return obj.get_bounding_box_per_env(num_envs)
+        return original_bbox_per_env(obj, num_envs)
+
+    has_het_sites = [
+        "isaaclab_arena.relations.bounding_box_helpers.has_heterogeneous_objects",
+        "isaaclab_arena.relations.object_placer.has_heterogeneous_objects",
+        "isaaclab_arena.relations.pooled_object_placer.has_heterogeneous_objects",
+    ]
+    bbox_per_env_sites = [
+        "isaaclab_arena.relations.bounding_box_helpers.get_bounding_box_per_env",
+        "isaaclab_arena.relations.object_placer.get_bounding_box_per_env",
+    ]
+    for site in has_het_sites:
+        monkeypatch.setattr(site, has_het_with_doubles)
+    for site in bbox_per_env_sites:
+        monkeypatch.setattr(site, bbox_per_env_with_doubles)
 
 
 # ---------------------------------------------------------------------------
@@ -448,13 +462,11 @@ def test_pooled_placer_heterogeneous_sample_with_replacement():
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
-    pool._layout_pools = {
-        env_id: [
+    for env_id in range(4):
+        pool._env_pools[env_id].layouts = [
             PlacementResult(success=True, positions={hetero: (float(env_id), 0.0, 0.0)}, final_loss=0.0, attempts=1)
         ]
-        for env_id in range(4)
-    }
-    pool._layout_cursors = {env_id: 0 for env_id in range(4)}
+        pool._env_pools[env_id].cursor = 0
     initial_remaining = pool.remaining
     samples = pool.sample_with_replacement(8)
     assert [sample.positions[hetero][0] for sample in samples] == [0.0, 1.0, 2.0, 3.0, 0.0, 1.0, 2.0, 3.0]
@@ -476,6 +488,53 @@ def test_pooled_placer_heterogeneous_sample_without_replacement_triggers_refill(
     assert len(draws) == 2, "Pool should refill and return requested layouts"
 
 
+def test_pooled_placer_seeded_refills_are_reproducible():
+    """Two seeded pools should produce the same layout sequence across a refill."""
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    placer_params = ObjectPlacerParams(
+        solver_params=solver_params,
+        placement_seed=11,
+        max_placement_attempts=2,
+    )
+
+    def _draw_sequence():
+        desk, hetero, _placer_params = _make_hetero_pool_objects()
+        pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=2, num_envs=2)
+        first_round = pool.sample_without_replacement(2)
+        refill_round = pool.sample_without_replacement(2)
+        return [result.positions[hetero] for result in first_round + refill_round]
+
+    assert _draw_sequence() == _draw_sequence()
+
+
+def test_pooled_placer_env_specific_fallbacks_are_reported(capsys):
+    """Env-specific best-loss fallbacks should be visible to callers."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=2, num_envs=2)
+    for env_pool in pool._env_pools:
+        env_pool.layouts = []
+        env_pool.cursor = 0
+
+    fallback_results = [
+        [
+            PlacementResult(
+                success=False,
+                positions={hetero: (float(cur_env), 0.0, 0.0)},
+                final_loss=1.0,
+                attempts=1,
+            )
+        ]
+        for cur_env in range(2)
+    ]
+
+    pool._store_env_matched_results(fallback_results, layouts_per_env=1)
+    captured = capsys.readouterr()
+
+    assert pool.had_fallbacks
+    assert "Falling back to best-loss layouts" in captured.out
+    assert [env_pool.available for env_pool in pool._env_pools] == [1, 1]
+
+
 def test_pooled_placer_reusable_layouts_report_complete_env_rounds():
     """Reusable layouts should still expose equal without-replacement capacity per env."""
     desk = _make_desk()
@@ -508,8 +567,9 @@ def test_pooled_placer_reusable_layouts_keep_partial_valid_results():
     placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=None)
 
     pool = PooledObjectPlacer(objects=[desk, box], placer_params=placer_params, pool_size=4, num_envs=4)
-    pool._layout_pools = {env_id: [] for env_id in range(4)}
-    pool._layout_cursors = {env_id: 0 for env_id in range(4)}
+    for env_pool in pool._env_pools:
+        env_pool.layouts = []
+        env_pool.cursor = 0
 
     layouts = [
         PlacementResult(success=True, positions={box: (float(i), 0.0, 0.0)}, final_loss=0.0, attempts=1)
@@ -517,7 +577,7 @@ def test_pooled_placer_reusable_layouts_keep_partial_valid_results():
     ]
     pool._store_reusable_results(layouts)
 
-    assert sum(len(pool._layout_pools[env_id]) for env_id in range(4)) == 3
+    assert sum(len(env_pool.layouts) for env_pool in pool._env_pools) == 3
     assert pool.remaining == 0
 
 
diff --git a/isaaclab_arena/tests/test_object_placer_init.py b/isaaclab_arena/tests/test_object_placer_init.py
index f44f3dfcb..e90072945 100644
--- a/isaaclab_arena/tests/test_object_placer_init.py
+++ b/isaaclab_arena/tests/test_object_placer_init.py
@@ -68,6 +68,33 @@ def test_on_init_z_places_bottom_at_parent_top():
     assert abs(child_bottom - expected_bottom) < 1e-6
 
 
+def test_on_init_uses_env_specific_parent_bbox_override():
+    """Object with On(anchor set) should initialize against that env's assigned bbox."""
+    table_set = DummyObject(
+        name="table_set",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(2.0, 2.0, 0.5)),
+    )
+    table_set.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)))
+    table_set.add_relation(IsAnchor())
+
+    small_table_bbox = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.3, 0.3, 0.1))
+    box_bbox = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.05, 0.05, 0.05))
+    box = DummyObject(name="box", bounding_box=box_bbox)
+    box.add_relation(On(table_set, clearance_m=0.02))
+
+    placer = ObjectPlacer(params=ObjectPlacerParams())
+    positions = placer._generate_initial_positions(
+        [table_set, box],
+        {table_set},
+        child_bboxes={table_set: small_table_bbox, box: box_bbox},
+    )
+
+    x, y, z = positions[box]
+    assert small_table_bbox.min_point[0, 0] <= x <= small_table_bbox.max_point[0, 0]
+    assert small_table_bbox.min_point[0, 1] <= y <= small_table_bbox.max_point[0, 1]
+    assert abs(z - (small_table_bbox.max_point[0, 2] + 0.02 - box_bbox.min_point[0, 2])) < 1e-6
+
+
 def test_on_init_clamps_to_center_when_child_wider_than_parent():
     """Object wider than its On parent in X/Y is clamped to parent center, not an invalid range."""
     desk = DummyObject(
diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py
index b41b3f645..375498760 100644
--- a/isaaclab_arena/tests/test_object_set.py
+++ b/isaaclab_arena/tests/test_object_set.py
@@ -94,6 +94,24 @@ def _test_object_set_default_variant_indices_follow_member_order(simulation_app)
     return True
 
 
+def _test_object_set_random_variant_indices_use_placement_seed(simulation_app):
+    """Random variant assignment should be repeatable with the same placement seed."""
+    from isaaclab_arena.assets.object_base import ObjectType
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+
+    def _assigned_indices():
+        can_a, can_b, _bbox_a, _bbox_b = _make_object_set_variants()
+        with (
+            patch("isaaclab_arena.assets.object_set.detect_object_type", return_value=ObjectType.RIGID),
+            patch("isaaclab_arena.assets.object_set.find_shallowest_rigid_body", return_value="/rigid"),
+        ):
+            obj_set = RigidObjectSet(name="seeded_cans", objects=[can_a, can_b], random_choice=True)
+            obj_set.assign_variants(num_envs=8, variant_seed=123)
+        return obj_set.variant_indices_by_env
+
+    return _assigned_indices() == _assigned_indices()
+
+
 def _test_object_set_rejects_variant_reassignment_with_different_num_envs(simulation_app):
     """Variant assignment should remain fixed for the original env count."""
     from isaaclab_arena.assets.object_base import ObjectType
@@ -211,15 +229,15 @@ def _test_empty_object_set(simulation_app):
 
 
 def _test_articulation_object_set(simulation_app):
+    from isaaclab_arena.assets.object_base import ObjectType
     from isaaclab_arena.assets.object_set import RigidObjectSet
-    from isaaclab_arena.assets.registries import AssetRegistry
 
-    asset_registry = AssetRegistry()
-    microwave = asset_registry.get_asset_by_name("microwave")()
+    can_a, can_b, _bbox_a, _bbox_b = _make_object_set_variants()
     try:
-        RigidObjectSet(name="articulation_object_set", objects=[microwave])
-    except Exception:
-        return True
+        with patch("isaaclab_arena.assets.object_set.detect_object_type", return_value=ObjectType.ARTICULATION):
+            RigidObjectSet(name="articulation_object_set", objects=[can_a, can_b])
+    except ValueError as exc:
+        return "contain only rigid objects" in str(exc)
     return False
 
 
@@ -473,6 +491,14 @@ def test_object_set_default_variant_indices_follow_member_order():
     assert result, f"Test {_test_object_set_default_variant_indices_follow_member_order.__name__} failed"
 
 
+def test_object_set_random_variant_indices_use_placement_seed():
+    result = run_simulation_app_function(
+        _test_object_set_random_variant_indices_use_placement_seed,
+        headless=HEADLESS,
+    )
+    assert result, f"Test {_test_object_set_random_variant_indices_use_placement_seed.__name__} failed"
+
+
 def test_object_set_rejects_variant_reassignment_with_different_num_envs():
     result = run_simulation_app_function(
         _test_object_set_rejects_variant_reassignment_with_different_num_envs,
diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py
index 947679191..d157c5d1d 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -274,7 +274,7 @@ def test_solve_and_place_objects_partial_reset_reusable_pool_consumes_only_reset
     assert available_before - available_after == len(env_ids)
 
 
-def test_solve_and_place_objects_writes_invalid_fallback_layout():
+def test_solve_and_place_objects_writes_invalid_fallback_layout(capsys):
     """Invalid fallback layouts should still be written, matching pool fallback behavior."""
 
     from isaaclab_arena.relations.placement_events import solve_and_place_objects
@@ -299,10 +299,53 @@ def sample_without_replacement(self, count: int) -> list[PlacementResult]:
             ]
 
     solve_and_place_objects(env, torch.tensor([0]), objects, InvalidPool())
+    captured = capsys.readouterr()
 
     assert set(env._assets) == {box1.name, box2.name}
     assert env._assets[box1.name].write_root_pose_to_sim.call_count == 1
     assert env._assets[box2.name].write_root_pose_to_sim.call_count == 1
+    assert "Writing best-loss fallback placement" in captured.out
+
+
+def test_solve_and_place_objects_partial_reset_env_indexed_uses_absolute_env_result():
+    """Env-indexed partial resets should write the result for each absolute env id."""
+
+    from isaaclab_arena.relations.placement_events import solve_and_place_objects
+    from isaaclab_arena.relations.placement_result import PlacementResult
+
+    desk, box1, box2 = _create_test_objects()
+    objects = [desk, box1, box2]
+    env = _make_mock_env(num_envs=4)
+
+    class EnvIndexedPool:
+        requires_env_indexed_layouts = True
+        num_envs = 4
+
+        def sample_without_replacement(self, count: int) -> list[PlacementResult]:
+            assert count == 4
+            return [
+                PlacementResult(
+                    success=True,
+                    positions={
+                        box1: (float(cur_env), 0.0, 0.0),
+                        box2: (float(cur_env), 1.0, 0.0),
+                    },
+                    final_loss=0.0,
+                    attempts=1,
+                )
+                for cur_env in range(4)
+            ]
+
+    solve_and_place_objects(env, torch.tensor([2]), objects, EnvIndexedPool())
+
+    box1_pose = env._assets[box1.name].write_root_pose_to_sim.call_args[0][0]
+    box2_pose = env._assets[box2.name].write_root_pose_to_sim.call_args[0][0]
+    box1_env_id = env._assets[box1.name].write_root_pose_to_sim.call_args.kwargs["env_ids"]
+    box2_env_id = env._assets[box2.name].write_root_pose_to_sim.call_args.kwargs["env_ids"]
+    assert box1_pose[0, 0].item() == 2.0
+    assert box2_pose[0, 0].item() == 2.0
+    assert box1_env_id.tolist() == [2]
+    assert box2_env_id.tolist() == [2]
 
 
 def test_solve_and_place_objects_asserts_env_indexed_pool_size_matches_scene():
@@ -318,7 +361,7 @@ class MismatchedEnvIndexedPool:
         requires_env_indexed_layouts = True
         num_envs = 1
 
-    with pytest.raises(AssertionError, match="scene has 2 env origins"):
+    with pytest.raises(ValueError, match="scene has 2 env origins"):
         solve_and_place_objects(env, torch.tensor([0]), objects, MismatchedEnvIndexedPool())
 
 

From f8fd190b7931f2c3ff643b8afac653866279905e Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 26 May 2026 15:57:37 -0700
Subject: [PATCH 40/46] remove list

---
 isaaclab_arena/assets/object_set.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 5f1c19598..45f902fc6 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -195,7 +195,7 @@ def _set_variant_indices_by_env(self, variant_indices_by_env: list[int]) -> None
         assert all(
             0 <= idx < n for idx in variant_indices_by_env
         ), f"RigidObjectSet '{self.name}' variant indices must be in [0, {n}); got {variant_indices_by_env}."
-        self.variant_indices_by_env = list(variant_indices_by_env)
+        self.variant_indices_by_env = variant_indices_by_env
         # During __init__, Object.object_cfg has not been built yet; _generate_rigid_cfg()
         # reads object_usd_paths after this assignment.
         spawn_cfg = self.object_cfg.spawn if getattr(self, "object_cfg", None) is not None else None

From 349c31dff606ec6464b84368fc653b2c0cb4c6fb Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 26 May 2026 19:02:04 -0700
Subject: [PATCH 41/46] edit comments

---
 isaaclab_arena/assets/object_set.py           | 22 ++++++++--------
 .../environments/arena_env_builder.py         |  6 ++---
 .../relations/bounding_box_helpers.py         |  8 +++---
 isaaclab_arena/relations/object_placer.py     | 14 +++++------
 .../relations/pooled_object_placer.py         | 25 +++++++++----------
 .../tests/test_heterogeneous_placement.py     |  8 +++---
 6 files changed, 41 insertions(+), 42 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 45f902fc6..af6eff747 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -98,7 +98,7 @@ def object_usd_paths(self) -> list[str]:
         """USD paths passed to MultiUsdFileCfg.
 
         Before assignment this is the member USD list. After assignment this
-        returns one USD path per environment based on ``variant_indices_by_env``.
+        returns one USD path per environment based on variant_indices_by_env.
         """
         if self.variant_indices_by_env is not None:
             return [self.member_usd_paths[idx] for idx in self.variant_indices_by_env]
@@ -107,9 +107,9 @@ def object_usd_paths(self) -> list[str]:
     def get_bounding_box(self) -> AxisAlignedBoundingBox:
         """Return one conservative local bbox for callers that cannot vary by env.
 
-        The returned bbox has shape ``(1, 3)`` and uses the member with the
+        The returned bbox has shape (1, 3) and uses the member with the
         greatest z-extent. Heterogeneous placement uses
-        ``get_bounding_box_per_env()`` after ``assign_variants()`` so each env
+        get_bounding_box_per_env() after assign_variants() so each env
         uses its actual variant geometry.
         """
         return max(self.objects, key=lambda obj: obj.get_bounding_box().size[0, 2].item()).get_bounding_box()
@@ -118,13 +118,13 @@ def assign_variants(self, num_envs: int, variant_seed: int | None = None) -> Non
         """Fix one member-variant index per environment.
 
         The assignment is fixed for the lifetime of the object set: subsequent
-        calls with the same ``num_envs`` are no-ops, and a call with a
-        different ``num_envs`` raises. When ``random_choice`` is True, each
+        calls with the same num_envs are no-ops, and a call with a
+        different num_envs raises. When random_choice is True, each
         env independently samples one variant; otherwise assignments repeat
         the member order across environments.
 
-        Callers invoke this once ``num_envs`` is known, before reading
-        ``variant_indices_by_env`` or ``get_bounding_box_per_env``.
+        Callers invoke this once num_envs is known, before reading
+        variant_indices_by_env or get_bounding_box_per_env.
 
         Args:
             num_envs: Number of environments to assign variants for.
@@ -145,15 +145,15 @@ def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         local bbox of the variant assigned to each env, enabling correct
         collision-free placement for heterogeneous scenes.
 
-        Requires ``assign_variants(num_envs)`` to have been called first. The
-        returned bbox has shape ``(num_envs, 3)``.
+        Requires assign_variants(num_envs) to have been called first. The
+        returned bbox has shape (num_envs, 3).
 
         Args:
             num_envs: Number of environments. Must match the assignment.
 
         Returns:
-            ``AxisAlignedBoundingBox`` with ``min_point`` / ``max_point`` of
-            shape ``(num_envs, 3)``.
+            AxisAlignedBoundingBox with min_point / max_point of
+            shape (num_envs, 3).
         """
         assert self.variant_indices_by_env is not None, (
             f"RigidObjectSet '{self.name}' has no variant assignment; "
diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py
index f0f4a5b8f..f11398e55 100644
--- a/isaaclab_arena/environments/arena_env_builder.py
+++ b/isaaclab_arena/environments/arena_env_builder.py
@@ -55,12 +55,12 @@ def _solve_relations(self) -> None:
            or by registering a pooled reset placement event
 
         Behaviour on reset depends on :attr:`ObjectPlacerParams.resolve_on_reset`
-        (overridable from CLI with ``--resolve_on_reset`` / ``--no-resolve_on_reset``):
+        (overridable from CLI with --resolve_on_reset / --no-resolve_on_reset):
 
         * **True** (default) — registers a reset event that draws a fresh layout
           from the pool for each resetting environment.
-        * **False** — applies one layout per environment via ``set_initial_pose``
-          so per-object reset events restore the same layout every time.
+        * **False** — applies one layout per environment so per-object reset
+          events restore the same layout every time.
         """
         objects_with_relations = self.arena_env.scene.get_objects_with_relations()
         self._placement_event_cfg = solve_and_apply_relation_placement(
diff --git a/isaaclab_arena/relations/bounding_box_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py
index 6f466b556..84527bc6f 100644
--- a/isaaclab_arena/relations/bounding_box_helpers.py
+++ b/isaaclab_arena/relations/bounding_box_helpers.py
@@ -17,8 +17,6 @@
 if TYPE_CHECKING:
     from isaaclab_arena.assets.object_base import ObjectBase
 
-VARIANT_SEED_STRIDE = 1_000_003
-
 
 def has_heterogeneous_objects(objects: list[ObjectBase]) -> bool:
     """Return whether placement must use env-specific object geometry."""
@@ -31,14 +29,16 @@ def assign_variants_for_envs(objects: list[ObjectBase], num_envs: int, placement
     """Assign per-env variants on every RigidObjectSet in the list.
 
     Placers call this once they know the real environment count, before
-    requesting per-env bounding boxes. Objects without variants are ignored.
+    requesting per-env bounding boxes. Non-RigidObjectSet objects are skipped.
+    Seeded assignments offset each set by its index so multiple sets do not
+    reuse the same random sequence.
     """
     from isaaclab_arena.assets.object_set import RigidObjectSet
 
     variant_set_idx = 0
     for obj in objects:
         if isinstance(obj, RigidObjectSet):
-            variant_seed = None if placement_seed is None else placement_seed + VARIANT_SEED_STRIDE * variant_set_idx
+            variant_seed = None if placement_seed is None else placement_seed + variant_set_idx
             obj.assign_variants(num_envs, variant_seed=variant_seed)
             variant_set_idx += 1
 
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index 6319bc1d4..8a2af1959 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -99,9 +99,9 @@ def place(
                 env must be solved against its assigned variant geometry.
 
         Returns:
-            ``PlacementResult`` for one produced layout, otherwise
-            ``MultiEnvPlacementResult``. A heterogeneous multi-env call produces
-            ``MultiEnvPlacementResult`` even when ``result_per_env`` is False.
+            PlacementResult for one produced layout, otherwise
+            MultiEnvPlacementResult. A heterogeneous multi-env call produces
+            MultiEnvPlacementResult even when result_per_env is False.
         """
         anchor_objects_set, generator = self._prepare_placement(objects)
 
@@ -147,10 +147,9 @@ def place_ranked_per_env(
     ) -> list[list[PlacementResult]]:
         """Return ranked placement candidates per env.
 
-        Use this for PooledObjectPlacer, where each env pool needs multiple
-        candidate layouts and will apply poses later when sampled. Use place()
-        for the normal public API that returns the selected placement result.
-        The return value has shape ``(num_envs, <=results_per_env)``: each
+        Use this for PooledObjectPlacer, where each env pool stores multiple
+        candidate layouts. Use place() for selected placement results.
+        The return value has shape (num_envs, <=results_per_env): each
         outer list entry corresponds to a real env, and each inner list is
         sorted with valid lower-loss layouts first.
         """
@@ -222,6 +221,7 @@ def _place_ranked(
                 " variants."
             )
 
+        # Variant assignment fixes the env-to-USD mapping before bbox expansion.
         assign_variants_for_envs(objects, num_envs, placement_seed=self.params.placement_seed)
         env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
         num_candidates = num_envs * candidates_per_env
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index a168ea611..164f347ff 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -41,8 +41,7 @@ def extend(self, layouts: list[PlacementResult]) -> None:
         self.layouts.extend(layouts)
 
     def next(self) -> PlacementResult:
-        if self.cursor >= len(self.layouts):
-            raise IndexError("No unread layouts remain in this env pool.")
+        assert self.cursor < len(self.layouts), "No unread layouts remain in this env pool."
         layout = self.layouts[self.cursor]
         self.cursor += 1
         return layout
@@ -57,14 +56,14 @@ class PooledObjectPlacer:
     consumed one at a time from the pooled queues.
 
     Strictly valid layouts are preferred. If no valid layout is available for
-    a solve batch, the best-loss solver result is kept as a visible fallback.
+    a solve batch, the best-loss solver result is kept as a fallback.
 
     The pool is refilled automatically when an env's queue runs out.
 
-    * ``sample_without_replacement(count)`` consumes ``count`` layouts. For
-      env-specific layouts, ``count`` must be a multiple of ``num_envs`` and
+    * sample_without_replacement(count) consumes count layouts. For
+      env-specific layouts, count must be a multiple of num_envs and
       may cover multiple complete env rounds.
-    * ``sample_with_replacement(count)`` is non-consuming. Env-specific layouts
+    * sample_with_replacement(count) is non-consuming. Env-specific layouts
       are sampled from matching env slots; reusable layouts are sampled IID from
       all stored layouts.
 
@@ -137,8 +136,8 @@ def _prepare_seeded_solve(self, num_candidates: int) -> None:
     def _solve_and_store(self, num_layouts: int) -> None:
         """Solve layouts in batches until every env has target_per_env unread layouts.
 
-        Each batch contributes (roughly) one round of layouts per env. The
-        outer loop is bounded by max_placement_attempts to avoid an
+        Each batch contributes one or more layout rounds per env. The outer
+        loop is bounded by max_placement_attempts to avoid an
         unbounded refill in pathological configurations.
         """
         self._discard_consumed_layouts()
@@ -215,8 +214,8 @@ def _solve_env_ranked_layouts(self, num_layouts: int) -> tuple[list[list[Placeme
         """Solve ranked layouts tied to each env's actual object geometry.
 
         Returns ranked candidate lists per real env so the pool can store
-        multiple layouts for each env without pretending extra candidate rows
-        are extra environments.
+        multiple layouts for each env without treating candidate rows as
+        environments.
         """
         layouts_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
         # ObjectPlacer seeds each candidate as placement_seed + candidate_idx.
@@ -270,12 +269,12 @@ def _store_env_matched_results(
     # ------------------------------------------------------------------
 
     def sample_without_replacement(self, count: int) -> list[PlacementResult]:
-        """Return the next ``count`` layouts.
+        """Return the next count layouts.
 
         Env-specific layouts are returned as complete rounds of
-        ``[env_0, env_1, ..., env_{num_envs-1}]`` so each result still maps
+        [env_0, env_1, ..., env_{num_envs-1}] so each result still maps
         to the absolute environment it was solved for. Reusable layouts consume
-        exactly ``count`` interchangeable entries.
+        exactly count interchangeable entries.
 
         Args:
             count: Number of layouts to return.
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index a35ab9b7e..46ff11cee 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -32,10 +32,10 @@
 
 @pytest.fixture(autouse=True)
 def _patch_bounding_box_helpers_for_test_doubles(monkeypatch):
-    """Let ``HeterogeneousDummyObject`` flow through the heterogeneous placement path.
+    """Let HeterogeneousDummyObject flow through the heterogeneous placement path.
 
-    Production dispatch uses ``isinstance(RigidObjectSet)``, but these tests use
-    lightweight ``DummyObject`` subclasses with ``get_bounding_box_per_env(...)``.
+    Production dispatch uses isinstance(RigidObjectSet), but these tests use
+    lightweight DummyObject subclasses with get_bounding_box_per_env(...).
     Patch every module that binds the helper by name at import time.
     """
     from isaaclab_arena.relations import bounding_box_helpers
@@ -508,7 +508,7 @@ def _draw_sequence():
 
 
 def test_pooled_placer_env_specific_fallbacks_are_reported(capsys):
-    """Env-specific best-loss fallbacks should be visible to callers."""
+    """Env-specific best-loss fallbacks should be reported to callers."""
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=2, num_envs=2)
     for env_pool in pool._env_pools:

From 86dd9afa4198fad600100358da5672771b51f4f4 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 27 May 2026 10:52:34 -0700
Subject: [PATCH 42/46] address comments

---
 isaaclab_arena/assets/object_set.py           |  38 +++--
 .../environments/relation_solver_interface.py |  27 ++--
 .../relations/bounding_box_helpers.py         |   7 +-
 isaaclab_arena/relations/object_placer.py     | 124 ++++++----------
 isaaclab_arena/relations/placement_events.py  |   9 +-
 .../relations/pooled_object_placer.py         |  78 +++++++---
 isaaclab_arena/relations/relation_solver.py   |   7 +-
 .../relations/relation_solver_state.py        |  11 +-
 .../tests/test_heterogeneous_placement.py     | 140 +++++++++++++++++-
 .../tests/test_no_collision_loss.py           |   6 +-
 .../tests/test_object_placer_init.py          |  26 +++-
 isaaclab_arena/tests/test_object_set.py       |   7 +-
 isaaclab_arena/tests/test_placement_events.py |  68 ++++++++-
 .../tests/test_validate_placement.py          |  28 ++--
 14 files changed, 388 insertions(+), 188 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index af6eff747..677038821 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -28,7 +28,6 @@ def __init__(
         prim_path: str | None = None,
         scale: tuple[float, float, float] = (1.0, 1.0, 1.0),
         random_choice: bool = False,
-        variant_indices_by_env: list[int] | None = None,
         initial_pose: Pose | None = None,
         **kwargs,
     ):
@@ -43,11 +42,10 @@ def __init__(
             random_choice: Whether to randomly choose an object from the object set to spawn in
                 each environment. If False, variants are assigned by repeating
                 the member order across environments.
-            variant_indices_by_env: Optional fixed variant index for each environment.
             initial_pose: The initial pose of the object from this object set.
         """
-        if len(objects) < 2:
-            raise ValueError(f"Object set {name} must contain at least 2 objects.")
+        if len(objects) < 1:
+            raise ValueError(f"Object set {name} must contain at least 1 object.")
         if not self._are_all_objects_type_rigid(objects):
             raise ValueError(f"Object set {name} must contain only rigid objects.")
 
@@ -77,9 +75,6 @@ def __init__(
         self.random_choice = random_choice
         self.variant_indices_by_env: list[int] | None = None
 
-        if variant_indices_by_env is not None:
-            self._set_variant_indices_by_env(variant_indices_by_env)
-
         if prim_path is None:
             prim_path = f"{{ENV_REGEX_NS}}/{name}"
 
@@ -105,7 +100,7 @@ def object_usd_paths(self) -> list[str]:
         return self.member_usd_paths
 
     def get_bounding_box(self) -> AxisAlignedBoundingBox:
-        """Return one conservative local bbox for callers that cannot vary by env.
+        """Return one local bbox for callers that cannot vary by env.
 
         The returned bbox has shape (1, 3) and uses the member with the
         greatest z-extent. Heterogeneous placement uses
@@ -117,11 +112,12 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox:
     def assign_variants(self, num_envs: int, variant_seed: int | None = None) -> None:
         """Fix one member-variant index per environment.
 
-        The assignment is fixed for the lifetime of the object set: subsequent
-        calls with the same num_envs are no-ops, and a call with a
-        different num_envs raises. When random_choice is True, each
-        env independently samples one variant; otherwise assignments repeat
-        the member order across environments.
+        The assignment is fixed for the lifetime of the object set so spawned
+        USDs and per-env bboxes stay aligned across placement refills.
+        Subsequent calls with the same num_envs are no-ops, and a call with a
+        different num_envs raises. When random_choice is True, each env
+        independently samples one variant; otherwise assignments repeat the
+        member order across environments.
 
         Callers invoke this once num_envs is known, before reading
         variant_indices_by_env or get_bounding_box_per_env.
@@ -130,13 +126,13 @@ def assign_variants(self, num_envs: int, variant_seed: int | None = None) -> Non
             num_envs: Number of environments to assign variants for.
             variant_seed: Optional seed used when random_choice=True.
         """
-        if self.variant_indices_by_env is not None:
+        if self.variant_indices_by_env is None:
+            self._set_variant_indices_by_env(self._generate_variant_indices(num_envs, variant_seed=variant_seed))
+        else:
             assert len(self.variant_indices_by_env) == num_envs, (
-                f"RigidObjectSet '{self.name}' has variant assignments for "
-                f"{len(self.variant_indices_by_env)} envs, got request for {num_envs}."
+                f"RigidObjectSet '{self.name}' got request for {num_envs} envs, "
+                f"but has variant assignments for {len(self.variant_indices_by_env)} envs."
             )
-            return
-        self._set_variant_indices_by_env(self._generate_variant_indices(num_envs, variant_seed=variant_seed))
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         """Return the local bbox for each env's assigned variant.
@@ -160,8 +156,8 @@ def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
             "call assign_variants(num_envs) before get_bounding_box_per_env()."
         )
         assert len(self.variant_indices_by_env) == num_envs, (
-            f"RigidObjectSet '{self.name}' assigned for "
-            f"{len(self.variant_indices_by_env)} envs, got request for {num_envs}."
+            f"RigidObjectSet '{self.name}' got request for {num_envs} envs, "
+            f"but is assigned for {len(self.variant_indices_by_env)} envs."
         )
         bounding_boxes = [obj.get_bounding_box() for obj in self.objects]
 
@@ -215,6 +211,8 @@ def _generate_rigid_cfg(self) -> RigidObjectCfg:
             prim_path=self.prim_path,
             spawn=sim_utils.MultiUsdFileCfg(
                 usd_path=self.object_usd_paths,
+                # Arena owns per-env variant assignment so bbox selection and
+                # spawned USDs stay aligned.
                 random_choice=False,
                 activate_contact_sensors=True,
             ),
diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py
index 5742fa6a5..4dd4f411f 100644
--- a/isaaclab_arena/environments/relation_solver_interface.py
+++ b/isaaclab_arena/environments/relation_solver_interface.py
@@ -111,21 +111,18 @@ def _apply_dynamic_spawn_pose(
     """Set initial spawn pose from one layout and return the reset placement event."""
     from isaaclab.managers import EventTermCfg
 
-    if placement_pool.requires_env_indexed_layouts:
-        print("Warning: Skipping static init_state seeding for env-indexed placement layouts.")
-    else:
-        layout = placement_pool.sample_with_replacement(1)[0]
-        for obj in objects:
-            if obj in anchor_objects_set:
-                continue
-            pos = layout.positions.get(obj)
-            if pos is None:
-                raise RuntimeError(f"Pool layout is missing object '{obj.name}'.")
-            object_cfg = getattr(obj, "object_cfg", None)
-            if object_cfg is None:
-                raise RuntimeError(f"Object '{obj.name}' must have object_cfg initialized before placement.")
-            object_cfg.init_state.pos = pos
-            object_cfg.init_state.rot = get_rotation_xyzw(obj)
+    layout = placement_pool.sample_with_replacement(1)[0]
+    for obj in objects:
+        if obj in anchor_objects_set:
+            continue
+        pos = layout.positions.get(obj)
+        if pos is None:
+            raise RuntimeError(f"Pool layout is missing object '{obj.name}'.")
+        object_cfg = getattr(obj, "object_cfg", None)
+        if object_cfg is None:
+            raise RuntimeError(f"Object '{obj.name}' must have object_cfg initialized before placement.")
+        object_cfg.init_state.pos = pos
+        object_cfg.init_state.rot = get_rotation_xyzw(obj)
 
     return EventTermCfg(
         func=solve_and_place_objects,
diff --git a/isaaclab_arena/relations/bounding_box_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py
index 84527bc6f..5a34f5ae2 100644
--- a/isaaclab_arena/relations/bounding_box_helpers.py
+++ b/isaaclab_arena/relations/bounding_box_helpers.py
@@ -12,6 +12,7 @@
 
 from typing import TYPE_CHECKING
 
+from isaaclab_arena.assets.object_set import RigidObjectSet
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
@@ -20,8 +21,6 @@
 
 def has_heterogeneous_objects(objects: list[ObjectBase]) -> bool:
     """Return whether placement must use env-specific object geometry."""
-    from isaaclab_arena.assets.object_set import RigidObjectSet
-
     return any(isinstance(obj, RigidObjectSet) for obj in objects)
 
 
@@ -33,8 +32,6 @@ def assign_variants_for_envs(objects: list[ObjectBase], num_envs: int, placement
     Seeded assignments offset each set by its index so multiple sets do not
     reuse the same random sequence.
     """
-    from isaaclab_arena.assets.object_set import RigidObjectSet
-
     variant_set_idx = 0
     for obj in objects:
         if isinstance(obj, RigidObjectSet):
@@ -49,8 +46,6 @@ def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBound
     RigidObjectSet delegates to its own get_bounding_box_per_env.
     All other objects broadcast their single bbox.
     """
-    from isaaclab_arena.assets.object_set import RigidObjectSet
-
     if isinstance(obj, RigidObjectSet):
         return obj.get_bounding_box_per_env(num_envs)
 
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index 8a2af1959..e3d78d067 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -69,16 +69,6 @@ def __init__(self, params: ObjectPlacerParams | None = None):
         self.params = params or ObjectPlacerParams()
         self._solver = RelationSolver(params=self.params.solver_params)
 
-    @staticmethod
-    def _resolve_bbox(
-        obj: ObjectBase,
-        overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None,
-    ) -> AxisAlignedBoundingBox:
-        """Return overrides[obj] if present, otherwise the object's default bbox."""
-        if overrides is not None and obj in overrides:
-            return overrides[obj]
-        return obj.get_bounding_box()
-
     def place(
         self,
         objects: list[ObjectBase],
@@ -223,11 +213,8 @@ def _place_ranked(
 
         # Variant assignment fixes the env-to-USD mapping before bbox expansion.
         assign_variants_for_envs(objects, num_envs, placement_seed=self.params.placement_seed)
-        env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
         num_candidates = num_envs * candidates_per_env
-        candidate_bboxes, per_env_bbox_overrides = self._build_candidate_bboxes(
-            env_bboxes, num_envs, candidates_per_env
-        )
+        candidate_bboxes, per_env_bboxes = self._build_candidate_bboxes(objects, num_envs, candidates_per_env)
 
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
         for candidate_idx in range(num_candidates):
@@ -235,9 +222,8 @@ def _place_ranked(
             if generator is not None:
                 assert self.params.placement_seed is not None
                 generator.manual_seed(self.params.placement_seed + candidate_idx)
-            env_child_bboxes = per_env_bbox_overrides[cur_env]
             initial_positions.append(
-                self._generate_initial_positions(objects, anchor_objects_set, generator, child_bboxes=env_child_bboxes)
+                self._generate_initial_positions(objects, anchor_objects_set, per_env_bboxes[cur_env], generator)
             )
 
         all_positions = self._solver.solve(objects, initial_positions, env_bboxes=candidate_bboxes)
@@ -247,12 +233,11 @@ def _place_ranked(
         candidates: list[PlacementCandidate] = []
         for candidate_idx in range(num_candidates):
             cur_env = candidate_idx // candidates_per_env
-            env_bbox_overrides = per_env_bbox_overrides[cur_env]
             candidates.append(
                 PlacementCandidate(
                     all_losses[candidate_idx],
                     all_positions[candidate_idx],
-                    self._validate_placement(all_positions[candidate_idx], bbox_overrides=env_bbox_overrides),
+                    self._validate_placement(all_positions[candidate_idx], per_env_bboxes[cur_env]),
                 )
             )
 
@@ -277,21 +262,22 @@ def _place_ranked(
 
     @staticmethod
     def _build_candidate_bboxes(
-        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox],
+        objects: list[ObjectBase],
         num_envs: int,
         candidates_per_env: int,
     ) -> tuple[dict[ObjectBase, AxisAlignedBoundingBox], list[dict[ObjectBase, AxisAlignedBoundingBox]]]:
         """Build solver bboxes with shape (num_envs * candidates_per_env, 3) and per-env views."""
-        per_env_bbox_overrides: list[dict[ObjectBase, AxisAlignedBoundingBox]] = [
-            {
-                obj: AxisAlignedBoundingBox(
+        env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
+
+        per_env_bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]] = []
+        for cur_env in range(num_envs):
+            cur_env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {}
+            for obj, bbox in env_bboxes.items():
+                cur_env_bboxes[obj] = AxisAlignedBoundingBox(
                     min_point=bbox.min_point[cur_env : cur_env + 1],
                     max_point=bbox.max_point[cur_env : cur_env + 1],
                 )
-                for obj, bbox in env_bboxes.items()
-            }
-            for cur_env in range(num_envs)
-        ]
+            per_env_bboxes.append(cur_env_bboxes)
 
         candidate_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {}
         for obj, bbox in env_bboxes.items():
@@ -300,7 +286,7 @@ def _build_candidate_bboxes(
                 max_point=bbox.max_point.repeat_interleave(candidates_per_env, dim=0),
             )
 
-        return candidate_bboxes, per_env_bbox_overrides
+        return candidate_bboxes, per_env_bboxes
 
     @staticmethod
     def _rank_candidates(
@@ -351,8 +337,8 @@ def _generate_initial_positions(
         self,
         objects: list[ObjectBase],
         anchor_objects: set[ObjectBase],
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox],
         generator: torch.Generator | None = None,
-        child_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> dict[ObjectBase, tuple[float, float, float]]:
         """Generate initial positions for all objects.
 
@@ -361,16 +347,15 @@ def _generate_initial_positions(
         anchor's center; the solver handles their placement from there.
 
         Args:
+            env_bboxes: Per-object bboxes for the current env, each with shape (1, 3).
             generator: Optional RNG generator for reproducible sampling. When None,
                 uses PyTorch's global RNG.
-            child_bboxes: Optional per-object bbox overrides with shape (1, 3).
-                Used to supply the correct env bbox for On-guided initialization.
 
         Returns:
             Dictionary mapping all objects to their starting positions.
         """
         first_anchor = next(obj for obj in objects if obj in anchor_objects)
-        anchor_bbox = self._get_world_bbox_for_init(first_anchor, child_bboxes)
+        anchor_bbox = self._get_world_bbox_for_init(first_anchor, env_bboxes)
 
         cx, cy, cz = float(anchor_bbox.center[0, 0]), float(anchor_bbox.center[0, 1]), float(anchor_bbox.center[0, 2])
 
@@ -378,16 +363,14 @@ def _generate_initial_positions(
         for obj in objects:
             if obj in anchor_objects:
                 initial_pose = obj.get_initial_pose()
-                if not isinstance(initial_pose, Pose):
-                    raise TypeError(
-                        f"Anchor object '{obj.name}' must have a fixed Pose before placement, got"
-                        f" {type(initial_pose).__name__}."
-                    )
+                assert isinstance(initial_pose, Pose), (
+                    f"Anchor object '{obj.name}' must have a fixed Pose before placement, got"
+                    f" {type(initial_pose).__name__}."
+                )
                 positions[obj] = initial_pose.position_xyz
             elif any(isinstance(r, On) for r in obj.get_relations()):
-                bbox_override = child_bboxes.get(obj) if child_bboxes else None
                 positions[obj] = self._compute_on_guided_position(
-                    obj, anchor_objects, anchor_bbox, generator, child_bbox=bbox_override, bbox_overrides=child_bboxes
+                    obj, anchor_objects, anchor_bbox, env_bboxes, generator
                 )
             else:
                 positions[obj] = (cx, cy, cz)
@@ -396,26 +379,20 @@ def _generate_initial_positions(
     @staticmethod
     def _get_world_bbox_for_init(
         obj: ObjectBase,
-        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox],
     ) -> AxisAlignedBoundingBox:
-        bbox_override = bbox_overrides.get(obj) if bbox_overrides else None
-        if bbox_override is None:
-            return obj.get_world_bounding_box()
-
         initial_pose = obj.get_initial_pose()
-        if not isinstance(initial_pose, Pose):
-            raise TypeError(
-                f"Object '{obj.name}' must have a fixed Pose to use an env-specific world bbox,"
-                f" got {type(initial_pose).__name__}."
-            )
-        return bbox_override.translated(initial_pose.position_xyz)
+        assert isinstance(
+            initial_pose, Pose
+        ), f"Object '{obj.name}' must have a fixed Pose to use its env bbox, got {type(initial_pose).__name__}."
+        return env_bboxes[obj].translated(initial_pose.position_xyz)
 
     def _get_on_parent_world_bbox(
         self,
         parent: ObjectBase,
         anchor_objects: set[ObjectBase],
         anchor_bbox: AxisAlignedBoundingBox,
-        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox],
     ) -> AxisAlignedBoundingBox:
         """Resolve the world bbox of an On relation's parent for initialization purposes.
 
@@ -427,10 +404,10 @@ def _get_on_parent_world_bbox(
         TODO(cvolk): Support full On-relation chains (e.g. spoon -> On(bowl) -> On(plate) -> On(table)).
         """
         if parent in anchor_objects:
-            return self._get_world_bbox_for_init(parent, bbox_overrides)
+            return self._get_world_bbox_for_init(parent, env_bboxes)
         for rel in parent.get_relations():
             if isinstance(rel, On) and rel.parent in anchor_objects:
-                return self._get_world_bbox_for_init(rel.parent, bbox_overrides)
+                return self._get_world_bbox_for_init(rel.parent, env_bboxes)
         return anchor_bbox
 
     def _compute_on_guided_position(
@@ -438,9 +415,8 @@ def _compute_on_guided_position(
         obj: ObjectBase,
         anchor_objects: set[ObjectBase],
         anchor_bbox: AxisAlignedBoundingBox,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox],
         generator: torch.Generator | None = None,
-        child_bbox: AxisAlignedBoundingBox | None = None,
-        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
     ) -> tuple[float, float, float]:
         """Compute an initial position for an object with an On relation.
 
@@ -448,19 +424,13 @@ def _compute_on_guided_position(
         so the solver starts from a valid region.
 
         Args:
+            env_bboxes: Per-object bboxes for the current env, each with shape (1, 3).
             generator: Optional RNG generator for reproducible sampling. When None,
                 uses PyTorch's global RNG.
-            child_bbox: Optional bbox override for the child object. When None,
-                obj.get_bounding_box() is used.
-            bbox_overrides: Optional per-object bbox overrides for resolving
-                heterogeneous On parents.
         """
         on_relation = next(r for r in obj.get_relations() if isinstance(r, On))
-        parent_bbox = self._get_on_parent_world_bbox(
-            on_relation.parent, anchor_objects, anchor_bbox, bbox_overrides=bbox_overrides
-        )
-        if child_bbox is None:
-            child_bbox = obj.get_bounding_box()
+        parent_bbox = self._get_on_parent_world_bbox(on_relation.parent, anchor_objects, anchor_bbox, env_bboxes)
+        child_bbox = env_bboxes[obj]
 
         x = self._sample_axis_position(
             parent_bbox.min_point[0, 0],
@@ -493,8 +463,8 @@ def _sample_axis_position(
         """Sample a child origin along one axis so the child's extent stays within the parent's extent.
 
         The valid range for the child origin is [parent_min - child_min, parent_max - child_max].
-        When low >= high, the child is wider than the parent on this axis — there's no position
-        where it fits completely, so we fall back to centering it over the parent.
+        When low >= high, the child is wider than the parent on this axis, so
+        return the parent center as a stable seed.
 
         Args:
             parent_min: Parent world-space min extent on this axis.
@@ -515,7 +485,7 @@ def _sample_axis_position(
     def _validate_on_relations(
         self,
         positions: dict[ObjectBase, tuple[float, float, float]],
-        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox],
     ) -> bool:
         """Validate each On relation; keep in sync with OnLossStrategy in relation_loss_strategies.py.
 
@@ -525,7 +495,7 @@ def _validate_on_relations(
 
         Args:
             positions: Solved positions for each object.
-            bbox_overrides: Optional per-object bbox overrides with shape (1, 3).
+            env_bboxes: Per-object bboxes for the current env, each with shape (1, 3).
         """
         for obj in positions:
             for rel in obj.get_relations():
@@ -534,8 +504,8 @@ def _validate_on_relations(
                 parent = rel.parent
                 if parent not in positions:
                     continue
-                child_bbox = self._resolve_bbox(obj, bbox_overrides)
-                parent_bbox = self._resolve_bbox(parent, bbox_overrides)
+                child_bbox = env_bboxes[obj]
+                parent_bbox = env_bboxes[parent]
                 child_world = child_bbox.translated(positions[obj])
                 parent_world = parent_bbox.translated(positions[parent])
                 if (
@@ -562,7 +532,7 @@ def _validate_on_relations(
     def _validate_no_overlap(
         self,
         positions: dict[ObjectBase, tuple[float, float, float]],
-        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox],
     ) -> bool:
         """Validate that no two objects overlap in 3D (axis-aligned bbox with margin).
 
@@ -573,7 +543,7 @@ def _validate_no_overlap(
 
         Args:
             positions: Solved positions for each object.
-            bbox_overrides: Optional per-object bbox overrides with shape (1, 3).
+            env_bboxes: Per-object bboxes for the current env, each with shape (1, 3).
         """
         on_pairs: set[tuple] = set()
         anchor_ids: set[int] = set()
@@ -600,8 +570,8 @@ def _validate_no_overlap(
                 if (id(a), id(b)) in on_pairs:
                     continue
 
-                a_bbox = self._resolve_bbox(a, bbox_overrides)
-                b_bbox = self._resolve_bbox(b, bbox_overrides)
+                a_bbox = env_bboxes[a]
+                b_bbox = env_bboxes[b]
                 a_world = a_bbox.translated(positions[a])
                 b_world = b_bbox.translated(positions[b])
 
@@ -614,20 +584,18 @@ def _validate_no_overlap(
     def _validate_placement(
         self,
         positions: dict[ObjectBase, tuple[float, float, float]],
-        bbox_overrides: dict[ObjectBase, AxisAlignedBoundingBox] | None = None,
+        env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox],
     ) -> bool:
         """Validate that no two objects overlap in 3D and On relations are satisfied.
 
         Args:
             positions: Dictionary mapping objects to their solved (x, y, z) positions.
-            bbox_overrides: Optional per-object bbox overrides with shape (1, 3).
+            env_bboxes: Per-object bboxes for the current env, each with shape (1, 3).
 
         Returns:
             True if no overlaps exist and On relations hold, False otherwise.
         """
-        return self._validate_no_overlap(positions, bbox_overrides) and self._validate_on_relations(
-            positions, bbox_overrides
-        )
+        return self._validate_no_overlap(positions, env_bboxes) and self._validate_on_relations(positions, env_bboxes)
 
     def _apply_positions(
         self,
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index 86056ab34..e1cd6629e 100644
--- a/isaaclab_arena/relations/placement_events.py
+++ b/isaaclab_arena/relations/placement_events.py
@@ -35,9 +35,9 @@ def solve_and_place_objects(
     """Coordinated reset event that draws layouts from the pool and writes poses.
 
     Registered as a single EventTermCfg(mode="reset"). Env-indexed pools
-    consume one complete scene round, then write only the resetting subset so
-    each result still matches its absolute env id. Reusable pools consume only
-    the number of resetting envs because those layouts are interchangeable.
+    consume one layout for each requested absolute env id. Reusable pools
+    consume only the number of resetting envs because those layouts are
+    interchangeable.
 
     Args:
         env: The Isaac Lab environment.
@@ -55,8 +55,7 @@ def solve_and_place_objects(
             raise ValueError(
                 f"Placement pool has {placement_pool.num_envs} envs, but scene has {num_scene_envs} env origins."
             )
-        all_results = placement_pool.sample_without_replacement(num_scene_envs)
-        results_by_env = {cur_env: all_results[cur_env] for cur_env in reset_env_ids}
+        results_by_env = placement_pool.sample_for_envs(reset_env_ids)
     else:
         reset_results = placement_pool.sample_without_replacement(len(reset_env_ids))
         results_by_env = dict(zip(reset_env_ids, reset_results))
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 164f347ff..b90be775d 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -51,12 +51,13 @@ class PooledObjectPlacer:
     """Object placer that maintains solved placement layouts.
 
     Storage is organized as one queue per environment. Env-specific layouts
-    are solved for a fixed env's object geometry and must be consumed in
-    complete env rounds. Reusable layouts are interchangeable and can be
-    consumed one at a time from the pooled queues.
+    are solved for fixed env geometry. sample_without_replacement consumes
+    complete env rounds; sample_for_envs consumes only the requested absolute
+    env ids. Reusable layouts are interchangeable and can be consumed one at a
+    time from the pooled queues.
 
-    Strictly valid layouts are preferred. If no valid layout is available for
-    a solve batch, the best-loss solver result is kept as a fallback.
+    Strictly valid layouts are preferred. On the final retry batch, best-loss
+    solver results may be kept as a fallback.
 
     The pool is refilled automatically when an env's queue runs out.
 
@@ -144,17 +145,23 @@ def _solve_and_store(self, num_layouts: int) -> None:
         target_per_env = max(1, (num_layouts + self._num_envs - 1) // self._num_envs)
         max_solve_batches = max(1, self._placer.params.max_placement_attempts)
 
-        for _ in range(max_solve_batches):
+        for batch_idx in range(max_solve_batches):
             max_missing = target_per_env - min(self._available_per_env())
             if max_missing <= 0:
                 return
 
             batch_size = max_missing * self._num_envs
+            allow_fallback = batch_idx == max_solve_batches - 1
             if self._uses_env_specific_bboxes:
                 ranked_results_per_env, layouts_per_env = self._solve_env_ranked_layouts(batch_size)
-                self._store_env_matched_results(ranked_results_per_env, layouts_per_env)
+                self._store_env_matched_results(
+                    ranked_results_per_env,
+                    layouts_per_env,
+                    allow_fallback=allow_fallback,
+                    target_per_env=target_per_env,
+                )
             else:
-                layouts = self._solve_reusable_layouts(batch_size)
+                layouts = self._solve_reusable_layouts(batch_size, allow_fallback=allow_fallback)
                 self._store_reusable_results(layouts)
 
             if min(self._available_per_env()) >= target_per_env:
@@ -165,12 +172,13 @@ def _solve_and_store(self, num_layouts: int) -> None:
             f"{max_solve_batches} solve batches. Available per env: {self._available_per_env()}."
         )
 
-    def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
+    def _solve_reusable_layouts(self, num_layouts: int, allow_fallback: bool = False) -> list[PlacementResult]:
         """Solve layouts that can be used by any env pool.
 
         Invalid candidates are discarded when at least one valid layout exists.
-        If no candidate passes strict validation, fall back to best-loss results
-        to match the pre-pool behavior used by existing environments.
+        If no candidate passes strict validation on the final retry batch, fall
+        back to best-loss results to match the pre-pool behavior used by
+        existing environments.
         """
         self._prepare_seeded_solve(num_layouts * self._placer.params.max_placement_attempts)
         with torch.inference_mode(False):
@@ -189,6 +197,9 @@ def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
         if valid_layouts:
             return valid_layouts
 
+        if not allow_fallback:
+            return []
+
         self._had_fallbacks = True
         print("Warning: No candidates passed strict validation. Accepting best-loss layouts as fallback.")
         return all_layouts
@@ -233,24 +244,32 @@ def _solve_env_ranked_layouts(self, num_layouts: int) -> tuple[list[list[Placeme
         return ranked_results_per_env, layouts_per_env
 
     def _store_env_matched_results(
-        self, ranked_results_per_env: list[list[PlacementResult]], layouts_per_env: int
+        self,
+        ranked_results_per_env: list[list[PlacementResult]],
+        layouts_per_env: int,
+        allow_fallback: bool = False,
+        target_per_env: int | None = None,
     ) -> None:
         """Store env-matched results into their corresponding pools.
 
-        Prefer successful layouts for each env. If a specific env has no valid
-        layouts in the batch, fall back to its best-loss results so existing
-        environments with imperfect validation can still run.
+        Prefer successful layouts for each env. If allow_fallback is set and a
+        specific env has no valid layouts, fall back to its best-loss results
+        so existing environments with imperfect validation can still run.
         """
         total_valid = 0
         fallback_envs = []
         for cur_env in range(self._num_envs):
             env_results = ranked_results_per_env[cur_env][:layouts_per_env]
             valid_results = [r for r in env_results if r.success]
+            missing = None if target_per_env is None else target_per_env - self._env_pools[cur_env].available
             if valid_results:
-                self._env_pools[cur_env].extend(valid_results)
                 total_valid += len(valid_results)
-            else:
-                self._env_pools[cur_env].extend(env_results)
+                if missing is None:
+                    self._env_pools[cur_env].extend(valid_results)
+                elif missing > 0:
+                    self._env_pools[cur_env].extend(valid_results[:missing])
+            elif allow_fallback and (missing is None or missing > 0):
+                self._env_pools[cur_env].extend(env_results if missing is None else env_results[:missing])
                 fallback_envs.append(cur_env)
                 self._had_fallbacks = True
 
@@ -308,6 +327,29 @@ def _sample_env_indexed_without_replacement(self, count: int) -> list[PlacementR
                 results.append(pool.next())
         return results
 
+    def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]:
+        """Consume one layout for each requested absolute env id."""
+        if not self._uses_env_specific_bboxes:
+            layouts = self._sample_reusable_without_replacement(len(env_ids))
+            return dict(zip(env_ids, layouts))
+
+        if any(env_id < 0 or env_id >= self._num_envs for env_id in env_ids):
+            raise ValueError(f"env_ids must be in [0, {self._num_envs}); got {env_ids}")
+
+        if any(self._env_pools[env_id].available < 1 for env_id in env_ids):
+            self._solve_and_store(max(self._pool_size, len(env_ids)))
+
+        results: dict[int, PlacementResult] = {}
+        for env_id in env_ids:
+            pool = self._env_pools[env_id]
+            if pool.available <= 0:
+                raise RuntimeError(
+                    f"Placement pool: env {env_id} has no more valid layouts. "
+                    "The solver is not producing enough valid placements."
+                )
+            results[env_id] = pool.next()
+        return results
+
     def _sample_reusable_without_replacement(self, count: int) -> list[PlacementResult]:
         """Consume exactly count interchangeable layouts."""
         if self._total_available() < count:
diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index ef4013261..e5a3b0a3e 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -237,10 +237,9 @@ def solve(
             initial_positions: List of dicts (one per env). Use a single-element list
                 for single-env placement.
             env_bboxes: Optional per-env bounding boxes keyed by object.
-                When provided, each AxisAlignedBoundingBox has shape
-                (batch, 3) so different batch rows can use different
-                geometry (heterogeneous placement). If None, every row
-                uses the object's default get_bounding_box().
+                ObjectPlacer always supplies these, with each
+                AxisAlignedBoundingBox shaped (batch, 3). Direct solver calls
+                may omit them to use each object's default get_bounding_box().
 
         Returns:
             List of dicts (one per env) mapping objects to their solved (x, y, z) positions.
diff --git a/isaaclab_arena/relations/relation_solver_state.py b/isaaclab_arena/relations/relation_solver_state.py
index fbb73219b..28b635f70 100644
--- a/isaaclab_arena/relations/relation_solver_state.py
+++ b/isaaclab_arena/relations/relation_solver_state.py
@@ -41,8 +41,9 @@ def __init__(
                 length > 1 = batched.
             device: Torch device for all tensors. Defaults to CPU.
             env_bboxes: Optional per-env bounding boxes keyed by object.
-                When provided, get_bbox(obj) returns the per-env bbox
-                (shape (batch, 3)) instead of the object's default.
+                ObjectPlacer always supplies these for placement solves. Direct
+                solver/debug calls may omit them to use each object's default
+                get_bounding_box().
         """
         assert len(initial_positions) >= 1, "initial_positions must contain at least one dict."
         anchor_objects = get_anchor_objects(objects)
@@ -150,11 +151,7 @@ def get_position(self, obj: ObjectBase) -> torch.Tensor:
         return self._optimizable_positions[:, opt_idx, :]
 
     def get_bbox(self, obj: ObjectBase) -> AxisAlignedBoundingBox:
-        """Return the local bounding box for obj, moved to the state's device.
-
-        Uses the per-env override from env_bboxes when available, otherwise
-        falls back to obj.get_bounding_box().
-        """
+        """Return the local bounding box for obj, moved to the state's device."""
         if self._env_bboxes is not None and obj in self._env_bboxes:
             return self._env_bboxes[obj].to(self._device)
         return obj.get_bounding_box().to(self._device)
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 46ff11cee..1121908fc 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -457,8 +457,32 @@ def test_pooled_placer_heterogeneous_sample_without_replacement_requires_complet
         pool.sample_without_replacement(2)
 
 
+def test_pooled_placer_sample_for_envs_consumes_only_requested_envs():
+    """sample_for_envs should advance only the requested absolute env pools."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
+
+    for env_id in range(4):
+        pool._env_pools[env_id].layouts = [
+            PlacementResult(
+                success=True,
+                positions={hetero: (float(env_id), 0.0, 0.0)},
+                final_loss=0.0,
+                attempts=1,
+            )
+        ]
+        pool._env_pools[env_id].cursor = 0
+
+    results = pool.sample_for_envs([2, 0])
+
+    assert list(results) == [2, 0]
+    assert results[2].positions[hetero][0] == 2.0
+    assert results[0].positions[hetero][0] == 0.0
+    assert [env_pool.available for env_pool in pool._env_pools] == [0, 1, 0, 1]
+
+
 def test_pooled_placer_heterogeneous_sample_with_replacement():
-    """sample_with_replacement should return per-variant layouts without consuming."""
+    """sample_with_replacement should return env-matched layouts without consuming."""
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=20, num_envs=4)
 
@@ -474,7 +498,7 @@ def test_pooled_placer_heterogeneous_sample_with_replacement():
 
 
 def test_pooled_placer_heterogeneous_sample_without_replacement_triggers_refill():
-    """Exhausting a variant sub-pool should trigger a refill."""
+    """Exhausting an env pool should trigger a refill."""
     desk, hetero, placer_params = _make_hetero_pool_objects()
     pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=4, num_envs=2)
 
@@ -514,6 +538,7 @@ def test_pooled_placer_env_specific_fallbacks_are_reported(capsys):
     for env_pool in pool._env_pools:
         env_pool.layouts = []
         env_pool.cursor = 0
+    pool._had_fallbacks = False
 
     fallback_results = [
         [
@@ -527,7 +552,7 @@ def test_pooled_placer_env_specific_fallbacks_are_reported(capsys):
         for cur_env in range(2)
     ]
 
-    pool._store_env_matched_results(fallback_results, layouts_per_env=1)
+    pool._store_env_matched_results(fallback_results, layouts_per_env=1, allow_fallback=True)
     captured = capsys.readouterr()
 
     assert pool.had_fallbacks
@@ -535,6 +560,115 @@ def test_pooled_placer_env_specific_fallbacks_are_reported(capsys):
     assert [env_pool.available for env_pool in pool._env_pools] == [1, 1]
 
 
+def test_pooled_placer_env_specific_fallbacks_wait_for_final_retry(capsys):
+    """Invalid env-specific candidates should not fill pools before fallback is allowed."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=2, num_envs=2)
+    for env_pool in pool._env_pools:
+        env_pool.layouts = []
+        env_pool.cursor = 0
+    pool._had_fallbacks = False
+
+    fallback_results = [
+        [
+            PlacementResult(
+                success=False,
+                positions={hetero: (float(cur_env), 0.0, 0.0)},
+                final_loss=1.0,
+                attempts=1,
+            )
+        ]
+        for cur_env in range(2)
+    ]
+
+    pool._store_env_matched_results(fallback_results, layouts_per_env=1)
+    captured = capsys.readouterr()
+
+    assert not pool.had_fallbacks
+    assert "Falling back to best-loss layouts" not in captured.out
+    assert [env_pool.available for env_pool in pool._env_pools] == [0, 0]
+
+
+def test_pooled_placer_env_specific_fallback_only_fills_short_env(capsys):
+    """Final-batch fallback should not overfill env pools that already met the target."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=2, num_envs=2)
+    for env_pool in pool._env_pools:
+        env_pool.layouts = []
+        env_pool.cursor = 0
+    pool._had_fallbacks = False
+
+    existing_layout = PlacementResult(
+        success=True,
+        positions={hetero: (10.0, 0.0, 0.0)},
+        final_loss=0.0,
+        attempts=1,
+    )
+    pool._env_pools[0].append(existing_layout)
+
+    fallback_results = [
+        [
+            PlacementResult(
+                success=False,
+                positions={hetero: (float(cur_env), 0.0, 0.0)},
+                final_loss=1.0,
+                attempts=1,
+            )
+        ]
+        for cur_env in range(2)
+    ]
+
+    pool._store_env_matched_results(
+        fallback_results,
+        layouts_per_env=1,
+        allow_fallback=True,
+        target_per_env=1,
+    )
+    captured = capsys.readouterr()
+
+    assert pool.had_fallbacks
+    assert "envs: [1]" in captured.out
+    assert [env_pool.available for env_pool in pool._env_pools] == [1, 1]
+    assert pool._env_pools[0].layouts == [existing_layout]
+    assert pool._env_pools[1].layouts[0] is fallback_results[1][0]
+
+
+def test_pooled_placer_env_specific_valid_results_only_fill_short_envs():
+    """Refills should not grow env pools that already met the target."""
+    desk, hetero, placer_params = _make_hetero_pool_objects()
+    pool = PooledObjectPlacer(objects=[desk, hetero], placer_params=placer_params, pool_size=2, num_envs=2)
+    for env_pool in pool._env_pools:
+        env_pool.layouts = []
+        env_pool.cursor = 0
+
+    existing_layout = PlacementResult(
+        success=True,
+        positions={hetero: (10.0, 0.0, 0.0)},
+        final_loss=0.0,
+        attempts=1,
+    )
+    pool._env_pools[0].append(existing_layout)
+
+    ranked_results = [
+        [
+            PlacementResult(
+                success=True,
+                positions={hetero: (float(cur_env), float(candidate_idx), 0.0)},
+                final_loss=0.0,
+                attempts=1,
+            )
+            for candidate_idx in range(2)
+        ]
+        for cur_env in range(2)
+    ]
+
+    pool._store_env_matched_results(ranked_results, layouts_per_env=2, target_per_env=1)
+
+    assert [env_pool.available for env_pool in pool._env_pools] == [1, 1]
+    assert pool._env_pools[0].layouts == [existing_layout]
+    assert pool._env_pools[1].layouts == [ranked_results[1][0]]
+
+
 def test_pooled_placer_reusable_layouts_report_complete_env_rounds():
     """Reusable layouts should still expose equal without-replacement capacity per env."""
     desk = _make_desk()
diff --git a/isaaclab_arena/tests/test_no_collision_loss.py b/isaaclab_arena/tests/test_no_collision_loss.py
index 0f2369157..3523ce410 100644
--- a/isaaclab_arena/tests/test_no_collision_loss.py
+++ b/isaaclab_arena/tests/test_no_collision_loss.py
@@ -266,9 +266,10 @@ def test_validation_accepts_on_parent_overlap():
     # the expanded table extends to 0.11, so they just touch. The On pair
     # should be skipped by validation.
     positions = {table: (0.0, 0.0, 0.0), box: (0.4, 0.4, 0.11)}
+    env_bboxes = {obj: obj.get_bounding_box() for obj in positions}
 
     placer = ObjectPlacer(ObjectPlacerParams())
-    assert placer._validate_no_overlap(positions)
+    assert placer._validate_no_overlap(positions, env_bboxes)
 
 
 def test_validation_rejects_non_anchor_overlap():
@@ -286,9 +287,10 @@ def test_validation_rejects_non_anchor_overlap():
 
     # Both boxes at nearly the same position -- they overlap
     positions = {table: (0.0, 0.0, 0.0), box_a: (0.3, 0.3, 0.11), box_b: (0.35, 0.35, 0.11)}
+    env_bboxes = {obj: obj.get_bounding_box() for obj in positions}
 
     placer = ObjectPlacer(ObjectPlacerParams())
-    assert not placer._validate_no_overlap(positions)
+    assert not placer._validate_no_overlap(positions, env_bboxes)
 
 
 def test_relation_solver_no_collision_same_inputs_reproducible():
diff --git a/isaaclab_arena/tests/test_object_placer_init.py b/isaaclab_arena/tests/test_object_placer_init.py
index e90072945..e26883e1e 100644
--- a/isaaclab_arena/tests/test_object_placer_init.py
+++ b/isaaclab_arena/tests/test_object_placer_init.py
@@ -24,6 +24,10 @@ def _make_desk():
     return desk
 
 
+def _env_bboxes(objects):
+    return {obj: obj.get_bounding_box() for obj in objects}
+
+
 def test_on_init_x_y_within_parent_footprint():
     """Object with On(anchor) is initialized with its bbox fully within parent's X/Y footprint."""
     desk = _make_desk()
@@ -34,7 +38,8 @@ def test_on_init_x_y_within_parent_footprint():
     box.add_relation(On(desk, clearance_m=0.01))
 
     placer = ObjectPlacer(params=ObjectPlacerParams())
-    positions = placer._generate_initial_positions([desk, box], {desk})
+    objects = [desk, box]
+    positions = placer._generate_initial_positions(objects, {desk}, _env_bboxes(objects))
 
     x, y, _ = positions[box]
     child_bbox = box.get_bounding_box()
@@ -57,7 +62,8 @@ def test_on_init_z_places_bottom_at_parent_top():
     box.add_relation(On(desk, clearance_m=clearance_m))
 
     placer = ObjectPlacer(params=ObjectPlacerParams())
-    positions = placer._generate_initial_positions([desk, box], {desk})
+    objects = [desk, box]
+    positions = placer._generate_initial_positions(objects, {desk}, _env_bboxes(objects))
 
     _, _, z = positions[box]
     child_bbox = box.get_bounding_box()
@@ -68,7 +74,7 @@ def test_on_init_z_places_bottom_at_parent_top():
     assert abs(child_bottom - expected_bottom) < 1e-6
 
 
-def test_on_init_uses_env_specific_parent_bbox_override():
+def test_on_init_uses_env_specific_parent_bbox():
     """Object with On(anchor set) should initialize against that env's assigned bbox."""
     table_set = DummyObject(
         name="table_set",
@@ -86,7 +92,7 @@ def test_on_init_uses_env_specific_parent_bbox_override():
     positions = placer._generate_initial_positions(
         [table_set, box],
         {table_set},
-        child_bboxes={table_set: small_table_bbox, box: box_bbox},
+        {table_set: small_table_bbox, box: box_bbox},
     )
 
     x, y, z = positions[box]
@@ -111,7 +117,8 @@ def test_on_init_clamps_to_center_when_child_wider_than_parent():
     big_box.add_relation(On(desk, clearance_m=0.0))
 
     placer = ObjectPlacer(params=ObjectPlacerParams())
-    positions = placer._generate_initial_positions([desk, big_box], {desk})
+    objects = [desk, big_box]
+    positions = placer._generate_initial_positions(objects, {desk}, _env_bboxes(objects))
 
     x, y, _ = positions[big_box]
     desk_world = desk.get_world_bounding_box()
@@ -132,7 +139,8 @@ def test_no_on_relation_initializes_at_anchor_center():
     box.add_relation(NextTo(desk, side=Side.POSITIVE_X, distance_m=0.05))
 
     placer = ObjectPlacer(params=ObjectPlacerParams())
-    positions = placer._generate_initial_positions([desk, box], {desk})
+    objects = [desk, box]
+    positions = placer._generate_initial_positions(objects, {desk}, _env_bboxes(objects))
 
     x, y, z = positions[box]
     desk_world = desk.get_world_bounding_box()
@@ -159,7 +167,8 @@ def test_on_non_anchor_parent_with_anchor_grandparent_uses_proxy():
     mug.add_relation(On(plate, clearance_m=0.0))
 
     placer = ObjectPlacer(params=ObjectPlacerParams())
-    positions = placer._generate_initial_positions([desk, plate, mug], {desk})
+    objects = [desk, plate, mug]
+    positions = placer._generate_initial_positions(objects, {desk}, _env_bboxes(objects))
 
     x, y, z = positions[mug]
     desk_world = desk.get_world_bounding_box()
@@ -188,7 +197,8 @@ def test_on_non_anchor_parent_without_on_uses_fallback_bbox():
     mug.add_relation(On(stand, clearance_m=0.0))
 
     placer = ObjectPlacer(params=ObjectPlacerParams())
-    positions = placer._generate_initial_positions([desk, stand, mug], {desk})
+    objects = [desk, stand, mug]
+    positions = placer._generate_initial_positions(objects, {desk}, _env_bboxes(objects))
 
     x, y, z = positions[mug]
     desk_world = desk.get_world_bounding_box()
diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py
index 375498760..7ba61a961 100644
--- a/isaaclab_arena/tests/test_object_set.py
+++ b/isaaclab_arena/tests/test_object_set.py
@@ -122,7 +122,8 @@ def _test_object_set_rejects_variant_reassignment_with_different_num_envs(simula
         patch("isaaclab_arena.assets.object_set.detect_object_type", return_value=ObjectType.RIGID),
         patch("isaaclab_arena.assets.object_set.find_shallowest_rigid_body", return_value="/rigid"),
     ):
-        obj_set = RigidObjectSet(name="preassigned_cans", objects=[can_a, can_b], variant_indices_by_env=[0, 1, 0])
+        obj_set = RigidObjectSet(name="assigned_cans", objects=[can_a, can_b])
+        obj_set.assign_variants(num_envs=3)
 
     try:
         obj_set.assign_variants(num_envs=4)
@@ -264,9 +265,7 @@ def _test_single_object_in_one_object_set(simulation_app):
         prim_path="{ENV_REGEX_NS}/kitchen/Cabinet_B_02",
         parent_asset=background,
     )
-    obj_set = RigidObjectSet(
-        name="single_object_set", objects=[cracker_box, cracker_box], prim_path=OBJECT_SET_1_PRIM_PATH
-    )
+    obj_set = RigidObjectSet(name="single_object_set", objects=[cracker_box], prim_path=OBJECT_SET_1_PRIM_PATH)
     obj_set.set_initial_pose(Pose(position_xyz=(0.1, 0.0, 0.1), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)))
     scene = Scene(assets=[background, obj_set])
     isaaclab_arena_environment = IsaacLabArenaEnvironment(
diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py
index d157c5d1d..964749a6a 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -320,11 +320,15 @@ def test_solve_and_place_objects_partial_reset_env_indexed_uses_absolute_env_res
     class EnvIndexedPool:
         requires_env_indexed_layouts = True
         num_envs = 4
+        requested_env_ids = None
 
         def sample_without_replacement(self, count: int) -> list[PlacementResult]:
-            assert count == 4
-            return [
-                PlacementResult(
+            raise AssertionError(f"partial reset should not consume a full env round, got count={count}")
+
+        def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]:
+            self.requested_env_ids = env_ids
+            return {
+                cur_env: PlacementResult(
                     success=True,
                     positions={
                         box1: (float(cur_env), 0.0, 0.0),
@@ -333,10 +337,11 @@ def sample_without_replacement(self, count: int) -> list[PlacementResult]:
                     final_loss=0.0,
                     attempts=1,
                 )
-                for cur_env in range(4)
-            ]
+                for cur_env in env_ids
+            }
 
-    solve_and_place_objects(env, torch.tensor([2]), objects, EnvIndexedPool())
+    pool = EnvIndexedPool()
+    solve_and_place_objects(env, torch.tensor([2]), objects, pool)
 
     box1_pose = env._assets[box1.name].write_root_pose_to_sim.call_args[0][0]
     box2_pose = env._assets[box2.name].write_root_pose_to_sim.call_args[0][0]
@@ -346,6 +351,7 @@ def sample_without_replacement(self, count: int) -> list[PlacementResult]:
     assert box2_pose[0, 0].item() == 2.0
     assert box1_env_id.tolist() == [2]
     assert box2_env_id.tolist() == [2]
+    assert pool.requested_env_ids == [2]
 
 
 def test_solve_and_place_objects_asserts_env_indexed_pool_size_matches_scene():
@@ -463,6 +469,56 @@ def test_resolve_on_reset_false_applies_pose_per_env():
             assert p.position_xyz is not None, f"Position should not be None for {obj.name}"
 
 
+def test_env_indexed_pool_seeds_init_state_before_reset_without_event():
+    """Env-indexed resolve-on-reset path should seed non-anchor initial poses."""
+    from types import SimpleNamespace
+
+    from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder
+    from isaaclab_arena.relations.placement_result import PlacementResult
+
+    class MinimalObject:
+        def __init__(self, name: str):
+            self.name = name
+            self.event_cfg = None
+            self.object_cfg = SimpleNamespace(init_state=SimpleNamespace(pos=(0.0, 0.0, 0.0), rot=(0.0, 0.0, 0.0, 1.0)))
+
+        def get_relations(self):
+            return []
+
+        def set_initial_pose(self, pose):
+            raise AssertionError("resolve_on_reset init seeding must not register per-object reset events")
+
+    class EnvIndexedPool:
+        requires_env_indexed_layouts = True
+        num_envs = 3
+        sample_count = None
+
+        def sample_with_replacement(self, count: int):
+            self.sample_count = count
+            assert count == 1
+            return [
+                PlacementResult(
+                    success=True,
+                    positions={box: (float(env_id), 0.0, 0.1)},
+                    final_loss=0.0,
+                    attempts=1,
+                )
+                for env_id in range(self.num_envs)
+            ]
+
+    anchor = MinimalObject("desk")
+    box = MinimalObject("box")
+    pool = EnvIndexedPool()
+    builder = ArenaEnvBuilder.__new__(ArenaEnvBuilder)
+
+    builder._set_init_state_from_pool([anchor, box], pool, {anchor})
+
+    assert pool.sample_count == 1
+    assert anchor.object_cfg.init_state.pos == (0.0, 0.0, 0.0)
+    assert box.object_cfg.init_state.pos == (0.0, 0.0, 0.1)
+    assert box.event_cfg is None
+
+
 def test_pooled_placer_falls_back_when_no_valid_layouts(capsys):
     """PooledObjectPlacer should keep best-loss fallback layouts when validation rejects all candidates."""
 
diff --git a/isaaclab_arena/tests/test_validate_placement.py b/isaaclab_arena/tests/test_validate_placement.py
index fc24835ee..484b19346 100644
--- a/isaaclab_arena/tests/test_validate_placement.py
+++ b/isaaclab_arena/tests/test_validate_placement.py
@@ -27,13 +27,17 @@ def _make_desk() -> DummyObject:
     )
 
 
+def _env_bboxes(positions: dict[DummyObject, tuple[float, float, float]]):
+    return {obj: obj.get_bounding_box() for obj in positions}
+
+
 def test_no_overlap_returns_true():
     """Test that two boxes far apart pass validation."""
     placer = ObjectPlacer(params=ObjectPlacerParams())
     a = _make_box("a")
     b = _make_box("b")
     positions = {a: (0.0, 0.0, 0.0), b: (1.0, 0.0, 0.0)}
-    assert placer._validate_placement(positions) is True
+    assert placer._validate_placement(positions, _env_bboxes(positions)) is True
 
 
 def test_overlapping_returns_false():
@@ -42,7 +46,7 @@ def test_overlapping_returns_false():
     a = _make_box("a")
     b = _make_box("b")
     positions = {a: (0.0, 0.0, 0.0), b: (0.0, 0.0, 0.0)}
-    assert placer._validate_placement(positions) is False
+    assert placer._validate_placement(positions, _env_bboxes(positions)) is False
 
 
 def test_partial_overlap_returns_false():
@@ -51,7 +55,7 @@ def test_partial_overlap_returns_false():
     a = _make_box("a", size=0.2)
     b = _make_box("b", size=0.2)
     positions = {a: (0.0, 0.0, 0.0), b: (0.1, 0.1, 0.0)}
-    assert placer._validate_placement(positions) is False
+    assert placer._validate_placement(positions, _env_bboxes(positions)) is False
 
 
 def test_separated_in_z_passes():
@@ -60,7 +64,7 @@ def test_separated_in_z_passes():
     a = _make_box("a")
     b = _make_box("b")
     positions = {a: (0.0, 0.0, 0.0), b: (0.0, 0.0, 5.0)}
-    assert placer._validate_placement(positions) is True
+    assert placer._validate_placement(positions, _env_bboxes(positions)) is True
 
 
 def test_object_on_surface_no_overlap():
@@ -70,7 +74,7 @@ def test_object_on_surface_no_overlap():
     box = _make_box("box", size=0.2)
     # Desk top at z=0.05; box at z=0.16 → box occupies z=[0.06, 0.26], clear of desk
     positions = {desk: (0.0, 0.0, 0.0), box: (0.0, 0.0, 0.16)}
-    assert placer._validate_placement(positions) is True
+    assert placer._validate_placement(positions, _env_bboxes(positions)) is True
 
 
 def test_colocated_siblings_overlap_rejected():
@@ -80,7 +84,7 @@ def test_colocated_siblings_overlap_rejected():
     a = _make_box("a", size=0.2)
     b = _make_box("b", size=0.2)
     positions = {desk: (0.0, 0.0, 0.0), a: (0.0, 0.0, 0.15), b: (0.0, 0.0, 0.15)}
-    assert placer._validate_placement(positions) is False
+    assert placer._validate_placement(positions, _env_bboxes(positions)) is False
 
 
 def test_on_relation_check_no_relation_returns_true():
@@ -89,7 +93,7 @@ def test_on_relation_check_no_relation_returns_true():
     a = _make_box("a")
     b = _make_box("b")
     positions = {a: (0.0, 0.0, 0.0), b: (1.0, 0.0, 0.0)}
-    assert placer._validate_on_relations(positions) is True
+    assert placer._validate_on_relations(positions, _env_bboxes(positions)) is True
 
 
 def test_on_relation_check_child_inside_xy_z_in_band_passes():
@@ -101,7 +105,7 @@ def test_on_relation_check_child_inside_xy_z_in_band_passes():
     box.add_relation(On(desk))  # clearance_m=0.01; desk top 0.05
     # Child bottom 0.06 (at upper bound); box half-height 0.1 → center z = 0.16.
     positions = {desk: (0.0, 0.0, 0.0), box: (0.0, 0.0, 0.16)}
-    assert placer._validate_on_relations(positions) is True
+    assert placer._validate_on_relations(positions, _env_bboxes(positions)) is True
 
 
 def test_validate_on_relations_child_z_above_clearance_fails():
@@ -113,7 +117,7 @@ def test_validate_on_relations_child_z_above_clearance_fails():
     box.add_relation(On(desk))  # clearance_m=0.01; desk top 0.05
     # Child bottom 1.0 is above band.
     positions = {desk: (0.0, 0.0, 0.0), box: (0.0, 0.0, 1.1)}
-    assert placer._validate_on_relations(positions) is False
+    assert placer._validate_on_relations(positions, _env_bboxes(positions)) is False
 
 
 def test_validate_on_relations_child_z_within_tolerance_above_clearance_passes():
@@ -125,7 +129,7 @@ def test_validate_on_relations_child_z_within_tolerance_above_clearance_passes()
     box.add_relation(On(desk))
     # Child bottom 0.063 → box center z = 0.063 + 0.1 = 0.163.
     positions = {desk: (0.0, 0.0, 0.0), box: (0.0, 0.0, 0.163)}
-    assert placer._validate_on_relations(positions) is True
+    assert placer._validate_on_relations(positions, _env_bboxes(positions)) is True
 
 
 def test_validate_on_relations_child_z_at_or_below_parent_top_fails():
@@ -136,7 +140,7 @@ def test_validate_on_relations_child_z_at_or_below_parent_top_fails():
     box = _make_box("box", size=0.2)
     box.add_relation(On(desk))  # clearance_m=0.01; desk top 0.05
     positions = {desk: (0.0, 0.0, 0.0), box: (0.0, 0.0, 0.15)}
-    assert placer._validate_on_relations(positions) is False
+    assert placer._validate_on_relations(positions, _env_bboxes(positions)) is False
 
 
 def test_on_relation_check_child_outside_xy_returns_false():
@@ -146,4 +150,4 @@ def test_on_relation_check_child_outside_xy_returns_false():
     box = _make_box("box", size=0.2)
     box.add_relation(On(desk))
     positions = {desk: (0.0, 0.0, 0.0), box: (10.0, 10.0, 0.1)}
-    assert placer._validate_on_relations(positions) is False
+    assert placer._validate_on_relations(positions, _env_bboxes(positions)) is False

From 9503c163ad10f6b8eb4464dc09c9b1a48c5242dd Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 27 May 2026 14:07:03 -0700
Subject: [PATCH 43/46] address comments

---
 isaaclab_arena/assets/object_set.py           | 18 ++++----
 .../relations/bounding_box_helpers.py         | 39 ++++++++++++++++++
 isaaclab_arena/relations/object_placer.py     | 41 ++++---------------
 .../tests/test_heterogeneous_placement.py     | 31 ++++++++++----
 isaaclab_arena/tests/test_object_set.py       | 29 ++++++++-----
 5 files changed, 98 insertions(+), 60 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index 677038821..c231e4021 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -114,10 +114,12 @@ def assign_variants(self, num_envs: int, variant_seed: int | None = None) -> Non
 
         The assignment is fixed for the lifetime of the object set so spawned
         USDs and per-env bboxes stay aligned across placement refills.
-        Subsequent calls with the same num_envs are no-ops, and a call with a
-        different num_envs raises. When random_choice is True, each env
+        Subsequent calls with the same num_envs are no-ops. A call with a
+        different num_envs regenerates with a warning. When random_choice is True, each env
         independently samples one variant; otherwise assignments repeat the
         member order across environments.
+        Regeneration is safe before the scene is spawned; afterwards, per-env
+        bboxes can desync from the spawned USDs.
 
         Callers invoke this once num_envs is known, before reading
         variant_indices_by_env or get_bounding_box_per_env.
@@ -126,13 +128,11 @@ def assign_variants(self, num_envs: int, variant_seed: int | None = None) -> Non
             num_envs: Number of environments to assign variants for.
             variant_seed: Optional seed used when random_choice=True.
         """
-        if self.variant_indices_by_env is None:
-            self._set_variant_indices_by_env(self._generate_variant_indices(num_envs, variant_seed=variant_seed))
-        else:
-            assert len(self.variant_indices_by_env) == num_envs, (
-                f"RigidObjectSet '{self.name}' got request for {num_envs} envs, "
-                f"but has variant assignments for {len(self.variant_indices_by_env)} envs."
-            )
+        if self.variant_indices_by_env is not None:
+            if len(self.variant_indices_by_env) == num_envs:
+                return
+            print(f"Warning: RigidObjectSet '{self.name}' regenerating variant assignments for {num_envs} envs.")
+        self._set_variant_indices_by_env(self._generate_variant_indices(num_envs, variant_seed=variant_seed))
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
         """Return the local bbox for each env's assigned variant.
diff --git a/isaaclab_arena/relations/bounding_box_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py
index 5a34f5ae2..cb8cd81fa 100644
--- a/isaaclab_arena/relations/bounding_box_helpers.py
+++ b/isaaclab_arena/relations/bounding_box_helpers.py
@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 from isaaclab_arena.assets.object_set import RigidObjectSet
@@ -54,3 +55,41 @@ def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBound
         min_point=bbox.min_point.expand(num_envs, 3),
         max_point=bbox.max_point.expand(num_envs, 3),
     )
+
+
+@dataclass
+class PerEnvBoundingBoxes:
+    """Per-env object bboxes with solver and validation formats."""
+
+    object_bboxes: dict[ObjectBase, AxisAlignedBoundingBox]
+    num_envs: int
+
+    def env_bboxes(self, env_id: int) -> dict[ObjectBase, AxisAlignedBoundingBox]:
+        """Return one-env bboxes for initialization and validation."""
+        return {
+            obj: AxisAlignedBoundingBox(
+                min_point=bbox.min_point[env_id : env_id + 1],
+                max_point=bbox.max_point[env_id : env_id + 1],
+            )
+            for obj, bbox in self.object_bboxes.items()
+        }
+
+    def all_env_bboxes(self) -> list[dict[ObjectBase, AxisAlignedBoundingBox]]:
+        """Return one-env bboxes for every env."""
+        return [self.env_bboxes(env_id) for env_id in range(self.num_envs)]
+
+    def solver_candidate_bboxes(self, candidates_per_env: int) -> dict[ObjectBase, AxisAlignedBoundingBox]:
+        """Return bboxes expanded to match candidate rows passed to the solver."""
+        return {
+            obj: AxisAlignedBoundingBox(
+                min_point=bbox.min_point.repeat_interleave(candidates_per_env, dim=0),
+                max_point=bbox.max_point.repeat_interleave(candidates_per_env, dim=0),
+            )
+            for obj, bbox in self.object_bboxes.items()
+        }
+
+
+def build_per_env_bounding_boxes(objects: list[ObjectBase], num_envs: int) -> PerEnvBoundingBoxes:
+    """Build per-env bboxes for each placement object."""
+    object_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
+    return PerEnvBoundingBoxes(object_bboxes=object_bboxes, num_envs=num_envs)
diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index e3d78d067..0108fceae 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -12,7 +12,7 @@
 
 from isaaclab_arena.relations.bounding_box_helpers import (
     assign_variants_for_envs,
-    get_bounding_box_per_env,
+    build_per_env_bounding_boxes,
     has_heterogeneous_objects,
 )
 from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams
@@ -214,7 +214,9 @@ def _place_ranked(
         # Variant assignment fixes the env-to-USD mapping before bbox expansion.
         assign_variants_for_envs(objects, num_envs, placement_seed=self.params.placement_seed)
         num_candidates = num_envs * candidates_per_env
-        candidate_bboxes, per_env_bboxes = self._build_candidate_bboxes(objects, num_envs, candidates_per_env)
+        env_bboxes = build_per_env_bounding_boxes(objects, num_envs)
+        candidate_bboxes = env_bboxes.solver_candidate_bboxes(candidates_per_env)
+        per_env_bboxes = env_bboxes.all_env_bboxes()
 
         initial_positions: list[dict[ObjectBase, tuple[float, float, float]]] = []
         for candidate_idx in range(num_candidates):
@@ -229,15 +231,18 @@ def _place_ranked(
         all_positions = self._solver.solve(objects, initial_positions, env_bboxes=candidate_bboxes)
         assert self._solver.last_loss_per_env is not None
         all_losses: list[float] = self._solver.last_loss_per_env.cpu().tolist()
+        all_validations = [
+            self._validate_placement(positions, per_env_bboxes[candidate_idx // candidates_per_env])
+            for candidate_idx, positions in enumerate(all_positions)
+        ]
 
         candidates: list[PlacementCandidate] = []
         for candidate_idx in range(num_candidates):
-            cur_env = candidate_idx // candidates_per_env
             candidates.append(
                 PlacementCandidate(
                     all_losses[candidate_idx],
                     all_positions[candidate_idx],
-                    self._validate_placement(all_positions[candidate_idx], per_env_bboxes[cur_env]),
+                    all_validations[candidate_idx],
                 )
             )
 
@@ -260,34 +265,6 @@ def _place_ranked(
 
         return ranked_results
 
-    @staticmethod
-    def _build_candidate_bboxes(
-        objects: list[ObjectBase],
-        num_envs: int,
-        candidates_per_env: int,
-    ) -> tuple[dict[ObjectBase, AxisAlignedBoundingBox], list[dict[ObjectBase, AxisAlignedBoundingBox]]]:
-        """Build solver bboxes with shape (num_envs * candidates_per_env, 3) and per-env views."""
-        env_bboxes = {obj: get_bounding_box_per_env(obj, num_envs) for obj in objects}
-
-        per_env_bboxes: list[dict[ObjectBase, AxisAlignedBoundingBox]] = []
-        for cur_env in range(num_envs):
-            cur_env_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {}
-            for obj, bbox in env_bboxes.items():
-                cur_env_bboxes[obj] = AxisAlignedBoundingBox(
-                    min_point=bbox.min_point[cur_env : cur_env + 1],
-                    max_point=bbox.max_point[cur_env : cur_env + 1],
-                )
-            per_env_bboxes.append(cur_env_bboxes)
-
-        candidate_bboxes: dict[ObjectBase, AxisAlignedBoundingBox] = {}
-        for obj, bbox in env_bboxes.items():
-            candidate_bboxes[obj] = AxisAlignedBoundingBox(
-                min_point=bbox.min_point.repeat_interleave(candidates_per_env, dim=0),
-                max_point=bbox.max_point.repeat_interleave(candidates_per_env, dim=0),
-            )
-
-        return candidate_bboxes, per_env_bboxes
-
     @staticmethod
     def _rank_candidates(
         candidates: list[PlacementCandidate],
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 1121908fc..b80605e84 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -14,7 +14,7 @@
 import pytest
 
 from isaaclab_arena.assets.dummy_object import DummyObject
-from isaaclab_arena.relations.bounding_box_helpers import get_bounding_box_per_env
+from isaaclab_arena.relations.bounding_box_helpers import build_per_env_bounding_boxes, get_bounding_box_per_env
 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, PlacementResult
@@ -36,7 +36,7 @@ def _patch_bounding_box_helpers_for_test_doubles(monkeypatch):
 
     Production dispatch uses isinstance(RigidObjectSet), but these tests use
     lightweight DummyObject subclasses with get_bounding_box_per_env(...).
-    Patch every module that binds the helper by name at import time.
+    Patch modules that bind the heterogeneous check by name at import time.
     """
     from isaaclab_arena.relations import bounding_box_helpers
 
@@ -56,14 +56,11 @@ def bbox_per_env_with_doubles(obj, num_envs):
         "isaaclab_arena.relations.object_placer.has_heterogeneous_objects",
         "isaaclab_arena.relations.pooled_object_placer.has_heterogeneous_objects",
     ]
-    bbox_per_env_sites = [
-        "isaaclab_arena.relations.bounding_box_helpers.get_bounding_box_per_env",
-        "isaaclab_arena.relations.object_placer.get_bounding_box_per_env",
-    ]
     for site in has_het_sites:
         monkeypatch.setattr(site, has_het_with_doubles)
-    for site in bbox_per_env_sites:
-        monkeypatch.setattr(site, bbox_per_env_with_doubles)
+    monkeypatch.setattr(
+        "isaaclab_arena.relations.bounding_box_helpers.get_bounding_box_per_env", bbox_per_env_with_doubles
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -120,6 +117,24 @@ def test_dummy_object_bbox_per_env_expands_single():
     assert torch.allclose(per_env.min_point[0], per_env.min_point[3])
 
 
+def test_per_env_bounding_boxes_formats_solver_and_env_views():
+    """PerEnvBoundingBoxes should expose solver and one-env bbox formats."""
+    obj = DummyObject(
+        name="box",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.3, 0.4)),
+    )
+
+    env_bboxes = build_per_env_bounding_boxes([obj], num_envs=3)
+    solver_bboxes = env_bboxes.solver_candidate_bboxes(candidates_per_env=2)
+    per_env_bboxes = env_bboxes.all_env_bboxes()
+
+    assert solver_bboxes[obj].min_point.shape == (6, 3)
+    assert solver_bboxes[obj].max_point.shape == (6, 3)
+    assert len(per_env_bboxes) == 3
+    assert per_env_bboxes[1][obj].min_point.shape == (1, 3)
+    assert torch.allclose(per_env_bboxes[1][obj].max_point[0], torch.tensor([0.2, 0.3, 0.4]))
+
+
 def test_heterogeneous_dummy_returns_different_bboxes():
     """HeterogeneousDummyObject should cycle through its member bboxes."""
 
diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py
index 7ba61a961..64aa54d22 100644
--- a/isaaclab_arena/tests/test_object_set.py
+++ b/isaaclab_arena/tests/test_object_set.py
@@ -112,8 +112,11 @@ def _assigned_indices():
     return _assigned_indices() == _assigned_indices()
 
 
-def _test_object_set_rejects_variant_reassignment_with_different_num_envs(simulation_app):
-    """Variant assignment should remain fixed for the original env count."""
+def _test_object_set_regenerates_variants_with_different_num_envs(simulation_app):
+    """Calling assign_variants with a different num_envs should regenerate indices."""
+    import io
+    from contextlib import redirect_stdout
+
     from isaaclab_arena.assets.object_base import ObjectType
     from isaaclab_arena.assets.object_set import RigidObjectSet
 
@@ -124,12 +127,16 @@ def _test_object_set_rejects_variant_reassignment_with_different_num_envs(simula
     ):
         obj_set = RigidObjectSet(name="assigned_cans", objects=[can_a, can_b])
         obj_set.assign_variants(num_envs=3)
-
-    try:
-        obj_set.assign_variants(num_envs=4)
-    except AssertionError:
-        return True
-    return False
+        assert obj_set.variant_indices_by_env is not None
+        assert len(obj_set.variant_indices_by_env) == 3
+
+        output = io.StringIO()
+        with redirect_stdout(output):
+            obj_set.assign_variants(num_envs=4)
+        assert obj_set.variant_indices_by_env is not None
+        assert len(obj_set.variant_indices_by_env) == 4
+        assert "regenerating variant assignments" in output.getvalue()
+    return True
 
 
 def _build_and_reset_env(simulation_app, scene_assets, env_name="object_set_test", task=None):
@@ -498,12 +505,12 @@ def test_object_set_random_variant_indices_use_placement_seed():
     assert result, f"Test {_test_object_set_random_variant_indices_use_placement_seed.__name__} failed"
 
 
-def test_object_set_rejects_variant_reassignment_with_different_num_envs():
+def test_object_set_regenerates_variants_with_different_num_envs():
     result = run_simulation_app_function(
-        _test_object_set_rejects_variant_reassignment_with_different_num_envs,
+        _test_object_set_regenerates_variants_with_different_num_envs,
         headless=HEADLESS,
     )
-    assert result, f"Test {_test_object_set_rejects_variant_reassignment_with_different_num_envs.__name__} failed"
+    assert result, f"Test {_test_object_set_regenerates_variants_with_different_num_envs.__name__} failed"
 
 
 def test_articulation_object_set():

From cab2c0260336a48ceb9c4e85f9a249b47d173a5b Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 27 May 2026 16:05:55 -0700
Subject: [PATCH 44/46] move rigidobject import

---
 isaaclab_arena/relations/bounding_box_helpers.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/isaaclab_arena/relations/bounding_box_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py
index cb8cd81fa..3ad602ea4 100644
--- a/isaaclab_arena/relations/bounding_box_helpers.py
+++ b/isaaclab_arena/relations/bounding_box_helpers.py
@@ -13,7 +13,6 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-from isaaclab_arena.assets.object_set import RigidObjectSet
 from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
 
 if TYPE_CHECKING:
@@ -22,6 +21,8 @@
 
 def has_heterogeneous_objects(objects: list[ObjectBase]) -> bool:
     """Return whether placement must use env-specific object geometry."""
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+
     return any(isinstance(obj, RigidObjectSet) for obj in objects)
 
 
@@ -33,6 +34,8 @@ def assign_variants_for_envs(objects: list[ObjectBase], num_envs: int, placement
     Seeded assignments offset each set by its index so multiple sets do not
     reuse the same random sequence.
     """
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+
     variant_set_idx = 0
     for obj in objects:
         if isinstance(obj, RigidObjectSet):
@@ -47,6 +50,8 @@ def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBound
     RigidObjectSet delegates to its own get_bounding_box_per_env.
     All other objects broadcast their single bbox.
     """
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+
     if isinstance(obj, RigidObjectSet):
         return obj.get_bounding_box_per_env(num_envs)
 

From 2bb1d51c0a3d2d6a3c498f5336c146c9719185cb Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Thu, 28 May 2026 13:27:18 -0700
Subject: [PATCH 45/46] refactor relation solver

---
 isaaclab_arena/assets/object_set.py           | 12 +---
 .../environments/relation_solver_interface.py | 35 +++++++----
 .../relations/bounding_box_helpers.py         | 12 +++-
 .../relations/pooled_object_placer.py         | 27 ++++----
 .../tests/test_heterogeneous_placement.py     | 53 ++++++++++++++++
 isaaclab_arena/tests/test_object_set.py       |  2 +-
 isaaclab_arena/tests/test_placement_events.py | 62 ++++++++++++++++++-
 .../tests/test_relation_solver_interface.py   |  6 +-
 8 files changed, 170 insertions(+), 39 deletions(-)

diff --git a/isaaclab_arena/assets/object_set.py b/isaaclab_arena/assets/object_set.py
index c231e4021..085cfc0b3 100644
--- a/isaaclab_arena/assets/object_set.py
+++ b/isaaclab_arena/assets/object_set.py
@@ -44,10 +44,8 @@ def __init__(
                 the member order across environments.
             initial_pose: The initial pose of the object from this object set.
         """
-        if len(objects) < 1:
-            raise ValueError(f"Object set {name} must contain at least 1 object.")
-        if not self._are_all_objects_type_rigid(objects):
-            raise ValueError(f"Object set {name} must contain only rigid objects.")
+        assert len(objects) >= 1, f"Object set {name} must contain at least 1 object."
+        assert self._are_all_objects_type_rigid(objects), f"Object set {name} must contain only rigid objects."
 
         # Isaac Lab support for MultiUsdFileCfg is limited. It applies the same scale and pose to all objects.
         # Furthermore it relies on the rigid body being at the root of the USD file, or at the same
@@ -135,11 +133,7 @@ def assign_variants(self, num_envs: int, variant_seed: int | None = None) -> Non
         self._set_variant_indices_by_env(self._generate_variant_indices(num_envs, variant_seed=variant_seed))
 
     def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
-        """Return the local bbox for each env's assigned variant.
-
-        Unlike the single-bbox compatibility fallback, this returns the real
-        local bbox of the variant assigned to each env, enabling correct
-        collision-free placement for heterogeneous scenes.
+        """Return each env's actual variant bbox.
 
         Requires assign_variants(num_envs) to have been called first. The
         returned bbox has shape (num_envs, 3).
diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py
index 4dd4f411f..f6f3a52b6 100644
--- a/isaaclab_arena/environments/relation_solver_interface.py
+++ b/isaaclab_arena/environments/relation_solver_interface.py
@@ -53,6 +53,8 @@ def solve_and_apply_relation_placement(
     if resolve_on_reset is not None:
         placer_params.resolve_on_reset = resolve_on_reset
 
+    # TODO(xinjieyao, 2026-05-22): Add joint object/embodiment placement once task-dependent
+    # reachability constraints are available. For now this always uses the object-only placer.
     placement_pool = PooledObjectPlacer(
         objects=objects,
         placer_params=placer_params,
@@ -82,8 +84,10 @@ def _apply_relation_placement_result(
 ) -> EventTermCfg | None:
     """Apply selected layouts to object spawn state and build reset event config."""
     anchor_objects_set = set(get_anchor_objects(objects))
+    # Prevent external pose-reset events from conflicting with relation-solved objects.
     _validate_no_conflicting_pose_reset_events(objects, anchor_objects_set)
 
+    # Anchor objects do not move, so no need to apply reset event.
     if anchor_objects_set == set(objects):
         return None
 
@@ -111,16 +115,16 @@ def _apply_dynamic_spawn_pose(
     """Set initial spawn pose from one layout and return the reset placement event."""
     from isaaclab.managers import EventTermCfg
 
+    # For env-indexed pools this seeds from env 0; the first reset overwrites with per-env layouts.
     layout = placement_pool.sample_with_replacement(1)[0]
     for obj in objects:
         if obj in anchor_objects_set:
             continue
         pos = layout.positions.get(obj)
         if pos is None:
-            raise RuntimeError(f"Pool layout is missing object '{obj.name}'.")
+            continue
         object_cfg = getattr(obj, "object_cfg", None)
-        if object_cfg is None:
-            raise RuntimeError(f"Object '{obj.name}' must have object_cfg initialized before placement.")
+        assert object_cfg is not None, f"Object '{obj.name}' must have object_cfg initialized before placement."
         object_cfg.init_state.pos = pos
         object_cfg.init_state.rot = get_rotation_xyzw(obj)
 
@@ -147,12 +151,20 @@ def _apply_static_initial_poses(
             continue
         rotation_xyzw = get_rotation_xyzw(obj)
         poses = []
+        missing_envs: list[int] = []
         for env_idx in range(num_envs):
             pos = layouts[env_idx].positions.get(obj)
             if pos is None:
-                raise RuntimeError(f"Placement layout for env {env_idx} is missing object '{obj.name}'.")
-            poses.append(Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw))
-        obj.set_initial_pose(PosePerEnv(poses=poses))
+                missing_envs.append(env_idx)
+            else:
+                poses.append(Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw))
+        if missing_envs:
+            print(
+                f"Warning: Object '{obj.name}' is missing positions in {len(missing_envs)} env(s) "
+                f"(env ids: {missing_envs}); skipping set_initial_pose for this object."
+            )
+        else:
+            obj.set_initial_pose(PosePerEnv(poses=poses))
 
 
 def _validate_no_conflicting_pose_reset_events(
@@ -161,9 +173,8 @@ def _validate_no_conflicting_pose_reset_events(
 ) -> None:
     """Reject conflicting explicit pose-reset events on relation-solved objects."""
     for obj in objects:
-        if obj not in anchor_objects_set and getattr(obj, "event_cfg", None) is not None:
-            raise RuntimeError(
-                f"Non-anchor object '{obj.name}' has an explicit pose-reset event. "
-                "Relational solving should not be combined with explicit setting of "
-                "poses on non-anchor objects."
-            )
+        assert not (obj not in anchor_objects_set and getattr(obj, "event_cfg", None) is not None), (
+            f"Non-anchor object '{obj.name}' has an explicit pose-reset event. "
+            "Relational solving should not be combined with explicit setting of "
+            "poses on non-anchor objects."
+        )
diff --git a/isaaclab_arena/relations/bounding_box_helpers.py b/isaaclab_arena/relations/bounding_box_helpers.py
index 3ad602ea4..cbdb9350d 100644
--- a/isaaclab_arena/relations/bounding_box_helpers.py
+++ b/isaaclab_arena/relations/bounding_box_helpers.py
@@ -62,13 +62,23 @@ def get_bounding_box_per_env(obj: ObjectBase, num_envs: int) -> AxisAlignedBound
     )
 
 
-@dataclass
+@dataclass(frozen=True)
 class PerEnvBoundingBoxes:
     """Per-env object bboxes with solver and validation formats."""
 
     object_bboxes: dict[ObjectBase, AxisAlignedBoundingBox]
     num_envs: int
 
+    def __post_init__(self) -> None:
+        assert self.num_envs >= 1, f"num_envs must be >= 1, got {self.num_envs}"
+        for obj, bbox in self.object_bboxes.items():
+            assert (
+                bbox.min_point.shape[0] == self.num_envs
+            ), f"Object '{obj.name}' bbox min_point has {bbox.min_point.shape[0]} envs, expected {self.num_envs}."
+            assert (
+                bbox.max_point.shape[0] == self.num_envs
+            ), f"Object '{obj.name}' bbox max_point has {bbox.max_point.shape[0]} envs, expected {self.num_envs}."
+
     def env_bboxes(self, env_id: int) -> dict[ObjectBase, AxisAlignedBoundingBox]:
         """Return one-env bboxes for initialization and validation."""
         return {
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index b90be775d..6fb9b0e40 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -83,14 +83,13 @@ def __init__(
         pool_size: int = 100,
         num_envs: int | None = None,
     ) -> None:
-        if pool_size < 1:
-            raise ValueError(f"pool_size must be >= 1, got {pool_size}")
+        assert pool_size >= 1, f"pool_size must be >= 1, got {pool_size}"
         self._uses_env_specific_bboxes = has_heterogeneous_objects(objects)
-        if self._uses_env_specific_bboxes and num_envs is None:
-            raise ValueError("num_envs is required when layouts use env-specific object variants.")
+        assert not (
+            self._uses_env_specific_bboxes and num_envs is None
+        ), "num_envs is required when layouts use env-specific object variants."
         self._num_envs = num_envs if num_envs is not None else 1
-        if self._num_envs < 1:
-            raise ValueError(f"num_envs must be >= 1, got {self._num_envs}")
+        assert self._num_envs >= 1, f"num_envs must be >= 1, got {self._num_envs}"
 
         self._objects = list(objects)
         # Pool construction ranks several candidate layouts per env and applies
@@ -177,8 +176,8 @@ def _solve_reusable_layouts(self, num_layouts: int, allow_fallback: bool = False
 
         Invalid candidates are discarded when at least one valid layout exists.
         If no candidate passes strict validation on the final retry batch, fall
-        back to best-loss results to match the pre-pool behavior used by
-        existing environments.
+        back to best-loss results so environments with imperfect validation can
+        still run.
         """
         self._prepare_seeded_solve(num_layouts * self._placer.params.max_placement_attempts)
         with torch.inference_mode(False):
@@ -263,11 +262,15 @@ def _store_env_matched_results(
             valid_results = [r for r in env_results if r.success]
             missing = None if target_per_env is None else target_per_env - self._env_pools[cur_env].available
             if valid_results:
-                total_valid += len(valid_results)
                 if missing is None:
+                    total_valid += len(valid_results)
                     self._env_pools[cur_env].extend(valid_results)
                 elif missing > 0:
-                    self._env_pools[cur_env].extend(valid_results[:missing])
+                    enqueued = valid_results[:missing]
+                    total_valid += len(enqueued)
+                    self._env_pools[cur_env].extend(enqueued)
+                else:
+                    total_valid += len(valid_results)
             elif allow_fallback and (missing is None or missing > 0):
                 self._env_pools[cur_env].extend(env_results if missing is None else env_results[:missing])
                 fallback_envs.append(cur_env)
@@ -276,7 +279,7 @@ def _store_env_matched_results(
         total_solved = sum(min(len(env_results), layouts_per_env) for env_results in ranked_results_per_env)
         if total_valid < total_solved:
             msg = (
-                f"Warning: Placement pool (env-specific bbox layouts) solved {total_solved} candidates,"
+                f"Placement pool (env-specific bbox layouts) solved {total_solved} candidates,"
                 f" {total_valid} valid, {total_solved - total_valid} failed validation"
             )
             if fallback_envs:
@@ -395,7 +398,7 @@ def sample_with_replacement(self, count: int) -> list[PlacementResult]:
 
         For env-specific layouts, slot i picks from env i % num_envs's pool
         so each result matches its absolute env. For reusable layouts, draws
-        are uniform IID from the full pool (preserving pre-heterogeneous behavior).
+        are uniform IID from the full pool.
         """
         if self._uses_env_specific_bboxes:
             results: list[PlacementResult] = []
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index b80605e84..577ee3ce7 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -865,3 +865,56 @@ def test_pooled_placer_per_env_pools_advance_in_complete_rounds():
     assert len(draws) == 4
     for d in draws:
         assert hetero in d.positions
+
+
+# ---------------------------------------------------------------------------
+# End-to-end with real RigidObjectSet
+# ---------------------------------------------------------------------------
+
+
+def test_real_rigid_object_set_through_pooled_placer():
+    """Real RigidObjectSet should flow through PooledObjectPlacer without monkey-patching.
+
+    This is an integration test that verifies the actual isinstance(RigidObjectSet)
+    dispatch in bounding_box_helpers.py triggers correctly, unlike the other tests
+    in this file that monkey-patch has_heterogeneous_objects.
+    """
+    from unittest.mock import patch
+
+    from isaaclab_arena.assets.object import Object
+    from isaaclab_arena.assets.object_base import ObjectType
+    from isaaclab_arena.assets.object_set import RigidObjectSet
+    from isaaclab_arena.relations.bounding_box_helpers import has_heterogeneous_objects
+
+    desk = _make_desk()
+
+    can_a = Object(name="can_a", object_type=ObjectType.RIGID, usd_path="/tmp/can_a.usd")
+    can_b = Object(name="can_b", object_type=ObjectType.RIGID, usd_path="/tmp/can_b.usd")
+    can_a.bounding_box = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.15))
+    can_b.bounding_box = AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.15, 0.15, 0.2))
+
+    with (
+        patch("isaaclab_arena.assets.object_set.detect_object_type", return_value=ObjectType.RIGID),
+        patch("isaaclab_arena.assets.object_set.find_shallowest_rigid_body", return_value="/rigid"),
+    ):
+        obj_set = RigidObjectSet(name="cans", objects=[can_a, can_b])
+
+    obj_set.add_relation(On(desk, clearance_m=0.01))
+
+    assert has_heterogeneous_objects([desk, obj_set])
+
+    num_envs = 4
+    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3, verbose=False)
+    placer_params = ObjectPlacerParams(solver_params=solver_params, placement_seed=42)
+
+    pool = PooledObjectPlacer(objects=[desk, obj_set], placer_params=placer_params, pool_size=20, num_envs=num_envs)
+
+    assert pool.requires_env_indexed_layouts
+    assert pool.remaining > 0
+
+    draws = pool.sample_without_replacement(num_envs)
+    assert len(draws) == num_envs
+    for draw in draws:
+        assert obj_set in draw.positions
+        z = draw.positions[obj_set][2]
+        assert abs(z - 0.11) < 0.05, f"z={z:.4f}, expected ~0.11"
diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py
index 64aa54d22..5014c66c5 100644
--- a/isaaclab_arena/tests/test_object_set.py
+++ b/isaaclab_arena/tests/test_object_set.py
@@ -244,7 +244,7 @@ def _test_articulation_object_set(simulation_app):
     try:
         with patch("isaaclab_arena.assets.object_set.detect_object_type", return_value=ObjectType.ARTICULATION):
             RigidObjectSet(name="articulation_object_set", objects=[can_a, can_b])
-    except ValueError as exc:
+    except AssertionError as exc:
         return "contain only rigid objects" in str(exc)
     return False
 
diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py
index 964749a6a..f4dc6ef3d 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -473,7 +473,7 @@ def test_env_indexed_pool_seeds_init_state_before_reset_without_event():
     """Env-indexed resolve-on-reset path should seed non-anchor initial poses."""
     from types import SimpleNamespace
 
-    from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder
+    from isaaclab_arena.environments.relation_solver_interface import _apply_dynamic_spawn_pose
     from isaaclab_arena.relations.placement_result import PlacementResult
 
     class MinimalObject:
@@ -509,9 +509,12 @@ def sample_with_replacement(self, count: int):
     anchor = MinimalObject("desk")
     box = MinimalObject("box")
     pool = EnvIndexedPool()
-    builder = ArenaEnvBuilder.__new__(ArenaEnvBuilder)
 
-    builder._set_init_state_from_pool([anchor, box], pool, {anchor})
+    _apply_dynamic_spawn_pose(
+        objects=[anchor, box],
+        placement_pool=pool,
+        anchor_objects_set={anchor},
+    )
 
     assert pool.sample_count == 1
     assert anchor.object_cfg.init_state.pos == (0.0, 0.0, 0.0)
@@ -519,6 +522,59 @@ def sample_with_replacement(self, count: int):
     assert box.event_cfg is None
 
 
+def test_env_indexed_static_poses_apply_per_env_positions():
+    """Static initial poses should apply per-env positions from env-indexed layouts."""
+    from isaaclab_arena.assets.dummy_object import DummyObject
+    from isaaclab_arena.environments.relation_solver_interface import _apply_static_initial_poses
+    from isaaclab_arena.relations.placement_result import PlacementResult
+    from isaaclab_arena.relations.relations import IsAnchor, On
+    from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox
+    from isaaclab_arena.utils.pose import Pose, PosePerEnv
+
+    desk = DummyObject(
+        name="desk",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)),
+    )
+    desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)))
+    desk.add_relation(IsAnchor())
+
+    box = DummyObject(
+        name="box",
+        bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)),
+    )
+    box.add_relation(On(desk, clearance_m=0.01))
+
+    num_envs = 3
+
+    class PerEnvPool:
+        requires_env_indexed_layouts = True
+        num_envs = 3
+
+        def sample_with_replacement(self, count: int):
+            return [
+                PlacementResult(
+                    success=True,
+                    positions={box: (0.1 * env_id, 0.2 * env_id, 0.11)},
+                    final_loss=0.0,
+                    attempts=1,
+                )
+                for env_id in range(count)
+            ]
+
+    _apply_static_initial_poses(
+        objects=[desk, box],
+        placement_pool=PerEnvPool(),
+        anchor_objects_set={desk},
+        num_envs=num_envs,
+    )
+
+    pose = box.get_initial_pose()
+    assert isinstance(pose, PosePerEnv)
+    assert len(pose.poses) == num_envs
+    for env_id in range(num_envs):
+        assert pose.poses[env_id].position_xyz == (0.1 * env_id, 0.2 * env_id, 0.11)
+
+
 def test_pooled_placer_falls_back_when_no_valid_layouts(capsys):
     """PooledObjectPlacer should keep best-loss fallback layouts when validation rejects all candidates."""
 
diff --git a/isaaclab_arena/tests/test_relation_solver_interface.py b/isaaclab_arena/tests/test_relation_solver_interface.py
index d60a881a2..a4a3e0054 100644
--- a/isaaclab_arena/tests/test_relation_solver_interface.py
+++ b/isaaclab_arena/tests/test_relation_solver_interface.py
@@ -32,6 +32,8 @@ def _make_box(name: str = "box"):
 
 
 class _FakePlacementPool:
+    requires_env_indexed_layouts = False
+
     def __init__(self, layouts) -> None:
         self._layouts = layouts
 
@@ -100,7 +102,7 @@ def test_dynamic_spawn_pose_skips_objects_missing_from_fallback_layout():
     assert box.get_initial_pose() is None
 
 
-def test_static_initial_poses_skip_object_when_any_layout_is_missing_position():
+def test_static_initial_poses_skip_object_when_any_layout_is_missing_position(capsys):
     from isaaclab_arena.environments.relation_solver_interface import _apply_static_initial_poses
     from isaaclab_arena.relations.placement_result import PlacementResult
     from isaaclab_arena.utils.pose import PosePerEnv
@@ -129,6 +131,8 @@ def test_static_initial_poses_skip_object_when_any_layout_is_missing_position():
         anchor_objects_set={desk},
         num_envs=2,
     )
+    captured = capsys.readouterr()
+    assert "missing_box" in captured.out
 
     assert missing_box.get_initial_pose() is None
     placed_box_initial_pose = placed_box.get_initial_pose()

From 483b0efc3cc3e5bc61a0745e4ba311b15bfab536 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Thu, 28 May 2026 14:30:03 -0700
Subject: [PATCH 46/46] unification of homo and heter mode

---
 isaaclab_arena/relations/object_placer.py     | 119 +++++-------------
 .../relations/object_placer_params.py         |   3 +-
 .../relations/pooled_object_placer.py         |  30 ++---
 isaaclab_arena/relations/relation_solver.py   |   4 +-
 .../tests/test_heterogeneous_placement.py     |  22 ++--
 .../test_object_placer_reproducibility.py     |  40 +-----
 6 files changed, 68 insertions(+), 150 deletions(-)

diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py
index 0108fceae..aff66023e 100644
--- a/isaaclab_arena/relations/object_placer.py
+++ b/isaaclab_arena/relations/object_placer.py
@@ -5,16 +5,11 @@
 
 from __future__ import annotations
 
-import math
 import torch
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING
 
-from isaaclab_arena.relations.bounding_box_helpers import (
-    assign_variants_for_envs,
-    build_per_env_bounding_boxes,
-    has_heterogeneous_objects,
-)
+from isaaclab_arena.relations.bounding_box_helpers import assign_variants_for_envs, build_per_env_bounding_boxes
 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
@@ -50,11 +45,11 @@ class ObjectPlacer:
     """High-level API for placing objects according to their spatial relations.
 
     Encapsulates the workflow of:
-    1. Random initialization of object positions
-    2. Running the RelationSolver
-    3. Validating the result
-    4. Retrying if necessary
-    5. Applying solved positions to objects
+    1. Random initialization of candidate positions per environment
+    2. Running the RelationSolver on all candidates in one batch
+    3. Validating each candidate
+    4. Ranking candidates per environment (valid first, then by loss)
+    5. Applying the best layout per environment to the objects
 
     Supports single-env (num_envs=1) and batched (num_envs>1) placement.
 
@@ -73,59 +68,41 @@ def place(
         self,
         objects: list[ObjectBase],
         num_envs: int = 1,
-        result_per_env: bool = True,
     ) -> PlacementResult | MultiEnvPlacementResult:
         """Place objects according to their spatial relations.
 
+        Every environment is solved against its own per-env bounding boxes and
+        receives its own best-ranked layout. Homogeneous objects share the same
+        bbox across envs; heterogeneous object sets use their assigned variant
+        geometry per env.
+
         Args:
             objects: List of objects to place. Must include at least one object
                 marked with IsAnchor() which serves as a fixed reference.
             num_envs: Number of environments. 1 for single-env; > 1 for batched
                 placement (one layout per env).
-            result_per_env: When True (default), each environment gets a distinct
-                layout. When False, homogeneous objects use a single best layout
-                for all environments. Heterogeneous object sets always return
-                one layout per env, even when this flag is False, because each
-                env must be solved against its assigned variant geometry.
 
         Returns:
-            PlacementResult for one produced layout, otherwise
-            MultiEnvPlacementResult. A heterogeneous multi-env call produces
-            MultiEnvPlacementResult even when result_per_env is False.
+            PlacementResult when num_envs == 1, otherwise a
+            MultiEnvPlacementResult with one layout per environment.
         """
         anchor_objects_set, generator = self._prepare_placement(objects)
-
-        uses_env_specific_bboxes = has_heterogeneous_objects(objects)
-        num_results = num_envs if result_per_env or uses_env_specific_bboxes else 1
         max_attempts = self.params.max_placement_attempts
-        if uses_env_specific_bboxes:
-            ranked_results = self._place_ranked(
-                objects,
-                anchor_objects_set,
-                num_envs,
-                candidates_per_env=max_attempts,
-                attempts_per_result=max_attempts,
-                ranking_mode="env_specific_geometry",
-                generator=generator,
-            )
-            results_per_env = [env_results[0] for env_results in ranked_results]
-        else:
-            ranked_results = self._place_ranked(
-                objects,
-                anchor_objects_set,
-                num_envs=num_results,
-                candidates_per_env=max_attempts,
-                attempts_per_result=max_attempts,
-                ranking_mode="shared_geometry",
-                generator=generator,
-            )
-            results_per_env = ranked_results[0][:num_results]
+        ranked_results_per_env = self._place_ranked(
+            objects,
+            anchor_objects_set,
+            num_envs,
+            candidates_per_env=max_attempts,
+            attempts_per_result=max_attempts,
+            generator=generator,
+        )
+        results_per_env = [env_results[0] for env_results in ranked_results_per_env]
 
-        positions_per_env = [r.positions for r in results_per_env]
         if self.params.apply_positions_to_objects:
+            positions_per_env = [r.positions for r in results_per_env]
             self._apply_positions(positions_per_env, anchor_objects_set)
 
-        if num_results == 1:
+        if num_envs == 1:
             return results_per_env[0]
         return MultiEnvPlacementResult(results=results_per_env)
 
@@ -139,7 +116,7 @@ def place_ranked_per_env(
 
         Use this for PooledObjectPlacer, where each env pool stores multiple
         candidate layouts. Use place() for selected placement results.
-        The return value has shape (num_envs, <=results_per_env): each
+        The return value has shape (num_envs, results_per_env): each
         outer list entry corresponds to a real env, and each inner list is
         sorted with valid lower-loss layouts first.
         """
@@ -152,7 +129,6 @@ def place_ranked_per_env(
             num_envs,
             candidates_per_env=max_attempts * results_per_env,
             attempts_per_result=max_attempts,
-            ranking_mode="env_specific_geometry",
             generator=generator,
         )
 
@@ -196,21 +172,14 @@ def _place_ranked(
         num_envs: int,
         candidates_per_env: int,
         attempts_per_result: int,
-        ranking_mode: Literal["shared_geometry", "env_specific_geometry"],
         generator: torch.Generator | None,
     ) -> list[list[PlacementResult]]:
-        """Solve and rank placement candidates.
+        """Solve and rank placement candidates per environment.
 
-        Shared-geometry ranking is used only when every candidate is
-        interchangeable. Env-specific-geometry ranking keeps candidates tied to
-        their env's assigned variants.
+        Each env is solved against its own per-env bounding boxes, and its
+        candidates are ranked independently (valid first, then by loss), so a
+        candidate is never compared against another env's geometry.
         """
-        if ranking_mode == "shared_geometry":
-            assert not has_heterogeneous_objects(objects), (
-                "Shared-geometry ranking requires homogeneous objects; candidates are not interchangeable across"
-                " variants."
-            )
-
         # Variant assignment fixes the env-to-USD mapping before bbox expansion.
         assign_variants_for_envs(objects, num_envs, placement_seed=self.params.placement_seed)
         num_candidates = num_envs * candidates_per_env
@@ -246,7 +215,7 @@ def _place_ranked(
                 )
             )
 
-        ranked_candidate_slices = self._rank_candidates(candidates, ranking_mode, num_envs, candidates_per_env)
+        ranked_candidate_slices = self._rank_candidates(candidates, num_envs, candidates_per_env)
         ranked_results = [
             [
                 PlacementResult(
@@ -261,21 +230,17 @@ def _place_ranked(
         ]
 
         if self.params.verbose:
-            self._print_ranked_summary(ranked_candidate_slices, ranking_mode, num_candidates, num_envs)
+            self._print_ranked_summary(ranked_candidate_slices, num_candidates, num_envs)
 
         return ranked_results
 
     @staticmethod
     def _rank_candidates(
         candidates: list[PlacementCandidate],
-        ranking_mode: Literal["shared_geometry", "env_specific_geometry"],
         num_envs: int,
         candidates_per_env: int,
     ) -> list[list[PlacementCandidate]]:
-        """Return one shared-geometry sorted slice or one sorted slice per env."""
-        if ranking_mode == "shared_geometry":
-            return [sorted(candidates, key=lambda candidate: (not candidate.is_valid, candidate.loss))]
-
+        """Return one loss-sorted candidate slice per env (valid candidates first)."""
         ranked_candidate_slices: list[list[PlacementCandidate]] = []
         for cur_env in range(num_envs):
             start = cur_env * candidates_per_env
@@ -288,27 +253,11 @@ def _rank_candidates(
     def _print_ranked_summary(
         self,
         ranked_candidate_slices: list[list[PlacementCandidate]],
-        ranking_mode: Literal["shared_geometry", "env_specific_geometry"],
         num_candidates: int,
         num_envs: int,
     ) -> None:
-        if ranking_mode == "shared_geometry":
-            candidates = ranked_candidate_slices[0]
-            total_valid = sum(1 for candidate in candidates if candidate.is_valid)
-            finite_losses = [candidate.loss for candidate in candidates if math.isfinite(candidate.loss)]
-            mean_loss = sum(finite_losses) / len(finite_losses) if finite_losses else float("inf")
-            n_valid = sum(1 for candidate in candidates[:num_envs] if candidate.is_valid)
-            print(
-                f"Solved {num_candidates} candidates in one batch: mean loss = {mean_loss:.6f},"
-                f" {total_valid} valid, selected best {num_envs} ({n_valid} valid)"
-            )
-            return
-
         n_valid = sum(1 for candidate_slice in ranked_candidate_slices if candidate_slice[0].is_valid)
-        print(
-            f"Solved {num_candidates} candidates in one batch (env-specific geometry): "
-            f"{n_valid}/{num_envs} env(s) valid"
-        )
+        print(f"Solved {num_candidates} candidates in one batch: {n_valid}/{num_envs} env(s) valid")
 
     def _generate_initial_positions(
         self,
diff --git a/isaaclab_arena/relations/object_placer_params.py b/isaaclab_arena/relations/object_placer_params.py
index a617f6f2e..d5ddb9d82 100644
--- a/isaaclab_arena/relations/object_placer_params.py
+++ b/isaaclab_arena/relations/object_placer_params.py
@@ -16,7 +16,8 @@ class ObjectPlacerParams:
     """Parameters for the underlying RelationSolver."""
 
     max_placement_attempts: int = 10
-    """Maximum number of placement attempts (random init + solve + validate) before failure."""
+    """Number of candidate layouts solved and ranked per result. Higher values raise the chance a valid
+    layout is found in the batched solve. Also bounds the refill batches in PooledObjectPlacer."""
 
     apply_positions_to_objects: bool = True
     """If True, automatically set solved positions on objects after placement."""
diff --git a/isaaclab_arena/relations/pooled_object_placer.py b/isaaclab_arena/relations/pooled_object_placer.py
index 6fb9b0e40..bbf43ecf5 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -64,6 +64,8 @@ class PooledObjectPlacer:
     * sample_without_replacement(count) consumes count layouts. For
       env-specific layouts, count must be a multiple of num_envs and
       may cover multiple complete env rounds.
+    * sample_for_envs(env_ids) consumes one layout for each requested
+      absolute env id (used for partial resets).
     * sample_with_replacement(count) is non-consuming. Env-specific layouts
       are sampled from matching env slots; reusable layouts are sampled IID from
       all stored layouts.
@@ -181,9 +183,9 @@ def _solve_reusable_layouts(self, num_layouts: int, allow_fallback: bool = False
         """
         self._prepare_seeded_solve(num_layouts * self._placer.params.max_placement_attempts)
         with torch.inference_mode(False):
-            result = self._placer.place(self._objects, num_envs=num_layouts, result_per_env=True)
+            result = self._placer.place(self._objects, num_envs=num_layouts)
 
-        # TODO(@zhx06): Simplify once ObjectPlacer.place() always returns MultiEnvPlacementResult.
+        # place() returns a single PlacementResult only when num_envs == 1.
         all_layouts = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
         valid_layouts = [layout for layout in all_layouts if layout.success]
 
@@ -246,33 +248,32 @@ def _store_env_matched_results(
         self,
         ranked_results_per_env: list[list[PlacementResult]],
         layouts_per_env: int,
+        target_per_env: int,
         allow_fallback: bool = False,
-        target_per_env: int | None = None,
     ) -> None:
         """Store env-matched results into their corresponding pools.
 
-        Prefer successful layouts for each env. If allow_fallback is set and a
-        specific env has no valid layouts, fall back to its best-loss results
-        so existing environments with imperfect validation can still run.
+        Each env is filled only up to target_per_env unread layouts, so envs
+        that already met the target are not overfilled. Successful layouts are
+        preferred; if allow_fallback is set and an env has no valid layouts,
+        fall back to its best-loss results so environments with imperfect
+        validation can still run.
         """
         total_valid = 0
         fallback_envs = []
         for cur_env in range(self._num_envs):
             env_results = ranked_results_per_env[cur_env][:layouts_per_env]
             valid_results = [r for r in env_results if r.success]
-            missing = None if target_per_env is None else target_per_env - self._env_pools[cur_env].available
+            missing = target_per_env - self._env_pools[cur_env].available
             if valid_results:
-                if missing is None:
-                    total_valid += len(valid_results)
-                    self._env_pools[cur_env].extend(valid_results)
-                elif missing > 0:
+                if missing > 0:
                     enqueued = valid_results[:missing]
                     total_valid += len(enqueued)
                     self._env_pools[cur_env].extend(enqueued)
                 else:
                     total_valid += len(valid_results)
-            elif allow_fallback and (missing is None or missing > 0):
-                self._env_pools[cur_env].extend(env_results if missing is None else env_results[:missing])
+            elif allow_fallback and missing > 0:
+                self._env_pools[cur_env].extend(env_results[:missing])
                 fallback_envs.append(cur_env)
                 self._had_fallbacks = True
 
@@ -400,6 +401,7 @@ def sample_with_replacement(self, count: int) -> list[PlacementResult]:
         so each result matches its absolute env. For reusable layouts, draws
         are uniform IID from the full pool.
         """
+        # With-replacement samples from all stored layouts, ignoring the consumption cursor.
         if self._uses_env_specific_bboxes:
             results: list[PlacementResult] = []
             for i in range(count):
@@ -423,7 +425,7 @@ def remaining(self) -> int:
 
     @property
     def pool_size(self) -> int:
-        """Number of layouts to solve per batch. When the pool runs low, it will solve at least this number of layouts so future samples can reuse the buffer."""
+        """Number of layouts solved per batch when the pool is refilled."""
         return self._pool_size
 
     @property
diff --git a/isaaclab_arena/relations/relation_solver.py b/isaaclab_arena/relations/relation_solver.py
index e5a3b0a3e..d11d456c6 100644
--- a/isaaclab_arena/relations/relation_solver.py
+++ b/isaaclab_arena/relations/relation_solver.py
@@ -138,9 +138,9 @@ def _compute_no_overlap_loss(
         state: RelationSolverState,
         debug: bool = False,
     ) -> torch.Tensor:
-        """Compute pairwise no-overlap loss for all non-anchor objects against all other objects.
+        """Compute pairwise no-overlap loss, skipping On-linked pairs.
 
-        Each unique pair is evaluated twice (once per direction):
+        Each unique non-On pair is evaluated twice (once per direction):
         - Non-anchor vs anchor: gradient flows to the non-anchor only.
         - Non-anchor vs non-anchor: both objects receive gradient by computing
           the loss in both directions with the other's position detached.
diff --git a/isaaclab_arena/tests/test_heterogeneous_placement.py b/isaaclab_arena/tests/test_heterogeneous_placement.py
index 577ee3ce7..e104020e3 100644
--- a/isaaclab_arena/tests/test_heterogeneous_placement.py
+++ b/isaaclab_arena/tests/test_heterogeneous_placement.py
@@ -53,7 +53,6 @@ def bbox_per_env_with_doubles(obj, num_envs):
 
     has_het_sites = [
         "isaaclab_arena.relations.bounding_box_helpers.has_heterogeneous_objects",
-        "isaaclab_arena.relations.object_placer.has_heterogeneous_objects",
         "isaaclab_arena.relations.pooled_object_placer.has_heterogeneous_objects",
     ]
     for site in has_het_sites:
@@ -244,7 +243,7 @@ def test_object_placer_heterogeneous_produces_per_env_results():
     )
 
     placer = ObjectPlacer(params=params)
-    result = placer.place(objects, num_envs=num_envs, result_per_env=True)
+    result = placer.place(objects, num_envs=num_envs)
 
     assert isinstance(result, MultiEnvPlacementResult)
     assert len(result.results) == num_envs
@@ -275,7 +274,7 @@ def test_object_placer_heterogeneous_z_height_matches_variant():
     )
 
     placer = ObjectPlacer(params=params)
-    result = placer.place(objects, num_envs=num_envs, result_per_env=True)
+    result = placer.place(objects, num_envs=num_envs)
 
     assert isinstance(result, MultiEnvPlacementResult)
     # Both envs should have solved z near the desk top + clearance (0.11).
@@ -321,7 +320,7 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
     )
 
     placer = ObjectPlacer(params=params)
-    result = placer.place(objects, num_envs=num_envs, result_per_env=True)
+    result = placer.place(objects, num_envs=num_envs)
 
     assert isinstance(result, MultiEnvPlacementResult)
     assert len(result.results) == num_envs
@@ -345,7 +344,7 @@ def test_mixed_heterogeneous_and_homogeneous_placement():
 
 
 def test_heterogeneous_placement_always_returns_per_env_results():
-    """Heterogeneous placement should not fall back to one shared approximate layout."""
+    """Heterogeneous placement returns one layout per env solved against its variant geometry."""
 
     desk, hetero, _placer_params = _make_hetero_pool_objects()
 
@@ -357,7 +356,7 @@ def test_heterogeneous_placement_always_returns_per_env_results():
     )
 
     placer = ObjectPlacer(params=params)
-    result = placer.place([desk, hetero], num_envs=4, result_per_env=False)
+    result = placer.place([desk, hetero], num_envs=4)
 
     assert isinstance(result, MultiEnvPlacementResult)
     assert len(result.results) == 4
@@ -398,8 +397,8 @@ def test_object_placer_place_ranked_per_env_returns_sorted_env_lists():
     assert hetero.get_initial_pose() is None
 
 
-def test_object_placer_homogeneous_path_returns_multi_env_result():
-    """When no heterogeneous objects exist, the homogeneous path is used."""
+def test_object_placer_homogeneous_objects_return_multi_env_result():
+    """Homogeneous objects return one layout per env (bboxes identical across envs)."""
 
     desk = _make_desk()
     box = DummyObject(
@@ -416,10 +415,11 @@ def test_object_placer_homogeneous_path_returns_multi_env_result():
     )
 
     placer = ObjectPlacer(params=params)
-    result = placer.place([desk, box], num_envs=2, result_per_env=True)
+    result = placer.place([desk, box], num_envs=2)
 
     assert isinstance(result, MultiEnvPlacementResult)
     assert len(result.results) == 2
+    assert all(r.success for r in result.results)
 
 
 # ---------------------------------------------------------------------------
@@ -567,7 +567,7 @@ def test_pooled_placer_env_specific_fallbacks_are_reported(capsys):
         for cur_env in range(2)
     ]
 
-    pool._store_env_matched_results(fallback_results, layouts_per_env=1, allow_fallback=True)
+    pool._store_env_matched_results(fallback_results, layouts_per_env=1, target_per_env=1, allow_fallback=True)
     captured = capsys.readouterr()
 
     assert pool.had_fallbacks
@@ -596,7 +596,7 @@ def test_pooled_placer_env_specific_fallbacks_wait_for_final_retry(capsys):
         for cur_env in range(2)
     ]
 
-    pool._store_env_matched_results(fallback_results, layouts_per_env=1)
+    pool._store_env_matched_results(fallback_results, layouts_per_env=1, target_per_env=1)
     captured = capsys.readouterr()
 
     assert not pool.had_fallbacks
diff --git a/isaaclab_arena/tests/test_object_placer_reproducibility.py b/isaaclab_arena/tests/test_object_placer_reproducibility.py
index 7db7b893f..bbcfdeea0 100644
--- a/isaaclab_arena/tests/test_object_placer_reproducibility.py
+++ b/isaaclab_arena/tests/test_object_placer_reproducibility.py
@@ -195,42 +195,8 @@ def test_relation_solver_multi_env_batched_positions():
             assert len(d[obj]) == 3
 
 
-def test_object_placer_result_per_env_false_returns_single_result():
-    """Test that place(num_envs>1, result_per_env=False) returns PlacementResult."""
-    num_envs = 4
-    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3)
-    desk, box1, box2 = _create_test_objects()
-    objects = [desk, box1, box2]
-    placer = ObjectPlacer(params=ObjectPlacerParams(placement_seed=42, solver_params=solver_params))
-    result = placer.place(objects, num_envs=num_envs, result_per_env=False)
-
-    assert isinstance(result, PlacementResult), "result_per_env=False should return PlacementResult"
-    assert not isinstance(result, MultiEnvPlacementResult)
-    assert box1 in result.positions
-    assert box2 in result.positions
-    assert len(result.positions[box1]) == 3
-    assert len(result.positions[box2]) == 3
-
-
-def test_object_placer_result_per_env_false_applies_pose_not_pose_per_env():
-    """Test that result_per_env=False sets a single Pose (not PosePerEnv) on each object."""
-    num_envs = 4
-    solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3)
-    desk, box1, box2 = _create_test_objects()
-    objects = [desk, box1, box2]
-    placer = ObjectPlacer(
-        params=ObjectPlacerParams(placement_seed=42, solver_params=solver_params, apply_positions_to_objects=True)
-    )
-    placer.place(objects, num_envs=num_envs, result_per_env=False)
-
-    for obj in [box1, box2]:
-        pose = obj.get_initial_pose()
-        assert isinstance(pose, Pose), f"{obj.name} should have a Pose, got {type(pose).__name__}"
-        assert not isinstance(pose, PosePerEnv)
-
-
-def test_object_placer_result_per_env_true_applies_pose_per_env():
-    """Test that result_per_env=True (default) sets PosePerEnv on each object when num_envs>1."""
+def test_object_placer_applies_pose_per_env():
+    """place(num_envs>1) sets PosePerEnv on each object."""
     num_envs = 4
     solver_params = RelationSolverParams(max_iters=200, convergence_threshold=1e-3)
     desk, box1, box2 = _create_test_objects()
@@ -238,7 +204,7 @@ def test_object_placer_result_per_env_true_applies_pose_per_env():
     placer = ObjectPlacer(
         params=ObjectPlacerParams(placement_seed=42, solver_params=solver_params, apply_positions_to_objects=True)
     )
-    placer.place(objects, num_envs=num_envs, result_per_env=True)
+    placer.place(objects, num_envs=num_envs)
 
     for obj in [box1, box2]:
         pose = obj.get_initial_pose()