From 78047acba0754a6e6432508ee108a976bb0dbd92 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 13 Apr 2026 14:34:02 -0700
Subject: [PATCH 01/24] 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 2ea9e3fdf..85e92b1ea 100644
--- a/isaaclab_arena/assets/object_base.py
+++ b/isaaclab_arena/assets/object_base.py
@@ -69,6 +69,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 75e2ea919fd40c90e0ae475902ddad49f2a27e4d Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 15 Apr 2026 14:06:57 -0700
Subject: [PATCH 02/24] 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 efecb5a74e136a87320ba38f6f1ffac91008db47 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 21 Apr 2026 09:38:19 -0700
Subject: [PATCH 03/24] 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/arena_env_builder.py         |   1 +
 isaaclab_arena/relations/placement_events.py  |   2 +-
 .../relations/pooled_object_placer.py         | 211 ++++++++++++++++--
 3 files changed, 194 insertions(+), 20 deletions(-)

diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py
index ff300095d..cfdb33391 100644
--- a/isaaclab_arena/environments/arena_env_builder.py
+++ b/isaaclab_arena/environments/arena_env_builder.py
@@ -92,6 +92,7 @@ def _solve_relations(self) -> None:
             objects=objects_with_relations,
             placer_params=placer_params,
             pool_size=pool_size,
+            num_envs=num_envs,
         )
 
         if placer_params.resolve_on_reset:
diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py
index aae6cc4c5..6ffc269af 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)
-    results_per_env = placement_pool.sample_without_replacement(num_reset_envs)
+    results_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 35fcadfda..4932aa3c2 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -24,8 +24,22 @@ class PooledObjectPlacer:
     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,68 @@ 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"Placement pool 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
+
+            self._solve_and_store(pool_size)
+            if not self._layouts:
+                raise RuntimeError(
+                    f"Placement pool 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,12 +133,9 @@ 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)
 
-        # TODO(@zhx06): Simplify once ObjectPlacer.place() always returns MultiEnvPlacementResult.
         all_results = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
         valid_results = [r for r in all_results if r.success]
 
@@ -93,20 +151,79 @@ def _solve_and_store(self, num_layouts: int) -> None:
             print("Warning: No candidates passed strict validation. Accepting best-loss layouts as fallback.")
             self._layouts.extend(all_results)
 
-    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)
+
         remaining = len(self._layouts) - self._next_idx
         if remaining < count:
-            self._solve_and_store(max(self._pool_size, count))  # solve a fresh batch
+            self._solve_and_store(max(self._pool_size, count))
 
         remaining = len(self._layouts) - self._next_idx
-        if remaining < count:  # still not enough after refill (solver producing too few valid layouts)
+        if remaining < count:
             raise RuntimeError(
                 f"Placement pool has {remaining} valid layouts but {count} were requested. "
                 "The solver is not producing enough valid placements."
@@ -116,15 +233,71 @@ def sample_without_replacement(self, count: int) -> list[PlacementResult]:
         self._next_idx += count
         return self._layouts[start : self._next_idx]
 
+    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

From ce6d8fddea21b3b87ae0b3cb382f7567ee8ca145 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Fri, 24 Apr 2026 10:41:44 -0700
Subject: [PATCH 04/24] 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 4932aa3c2..72fe2b34c 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)
@@ -189,8 +189,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 74043aa45fbd28893bbc157daba08bdb8afed46e Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Fri, 24 Apr 2026 11:10:10 -0700
Subject: [PATCH 05/24] 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 72fe2b34c..433bef79c 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:
@@ -96,25 +92,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
     # ------------------------------------------------------------------
@@ -152,48 +132,73 @@ def _solve_and_store(self, num_layouts: int) -> None:
             self._layouts.extend(all_results)
 
     # ------------------------------------------------------------------
-    # 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
@@ -208,7 +213,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.
@@ -239,6 +244,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):
@@ -247,28 +253,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
 
@@ -276,8 +286,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.
@@ -287,11 +296,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
 
@@ -299,8 +308,8 @@ 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
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 4054b088499440ea5975d226d6c27a62324456f3 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 4 May 2026 12:55:10 -0700
Subject: [PATCH 06/24] 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 ac8415e0832712e1c7a7ef5851f1dfc5c3f7801e Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 4 May 2026 16:39:22 -0700
Subject: [PATCH 07/24] 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 433bef79c..3d3e9d15c 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -144,55 +144,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
+
+        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)
+            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 _ in range(num_rounds):
-            with torch.inference_mode(False):
-                result = self._placer.place(self._objects, num_envs=self._num_envs, result_per_env=True)
+        with torch.inference_mode(False):
+            result = self._placer.place(
+                self._objects,
+                num_envs=total_layouts,
+                result_per_env=True,
+                env_bboxes=tiled_bboxes,
+            )
 
-            round_results = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
-            all_round_results.append(round_results)
-            total_solved += len(round_results)
+        all_results = result.results if isinstance(result, MultiEnvPlacementResult) else [result]
 
-            for env_id, r in enumerate(round_results):
-                if r.success:
-                    self._layout_pools[env_id].append(r)
-                    total_valid += 1
+        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 ba2a96efe0401e408104cfd447e8bdbb98cecf63 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 15:11:08 -0700
Subject: [PATCH 08/24] 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 0f4a9acfbbe33dee822666d60aeebfe8ff82e9fd Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 15:11:29 -0700
Subject: [PATCH 09/24] 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 a291a1bf131bb33d377d67b0aa8810728c0d2df5 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 15:37:43 -0700
Subject: [PATCH 10/24] 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 cfdb33391..f15980f04 100644
--- a/isaaclab_arena/environments/arena_env_builder.py
+++ b/isaaclab_arena/environments/arena_env_builder.py
@@ -388,7 +388,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 aadd8bd25c4ecad639779941d8ae29765ca7e110 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 15:49:18 -0700
Subject: [PATCH 11/24] 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 f0803c2a9504a5ae5fd9bf4cd772a421b8e9e59e Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 16:01:24 -0700
Subject: [PATCH 12/24] 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 944405b7673675a9e68f8526323b38f014214e30 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 16:18:23 -0700
Subject: [PATCH 13/24] 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 f15980f04..cdeca8327 100644
--- a/isaaclab_arena/environments/arena_env_builder.py
+++ b/isaaclab_arena/environments/arena_env_builder.py
@@ -388,7 +388,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 6709cfcb31a89ddf3db7a640d48d96b745a882e6 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 18:11:40 -0700
Subject: [PATCH 14/24] 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 0f603bee68ea5c67aa630ea105696825836e6cf5 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 5 May 2026 18:27:30 -0700
Subject: [PATCH 15/24] 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 f92fd3e5e6e1b09898f661190543532d0edab343 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 12:10:31 -0700
Subject: [PATCH 16/24] 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 379567fb0..c26430c0e 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
@@ -138,6 +139,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 cdeca8327..cfdb33391 100644
--- a/isaaclab_arena/environments/arena_env_builder.py
+++ b/isaaclab_arena/environments/arena_env_builder.py
@@ -388,30 +388,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 69de18deeae784b2128c5ad9649ecb5de1887df8 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 14:40:50 -0700
Subject: [PATCH 17/24] clean pool strategies

---
 .../isaaclab_arena_environment.py             |   4 +-
 isaaclab_arena/relations/placement_events.py  |  12 +-
 .../relations/pooled_object_placer.py         | 283 ++++++++----------
 .../tests/test_heterogeneous_placement.py     | 125 +++++---
 isaaclab_arena/tests/test_placement_events.py |   4 +-
 5 files changed, 224 insertions(+), 204 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 6ffc269af..7a53f0598 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 pooled placer 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)
-    results_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 = results_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 3d3e9d15c..d30126c1e 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -24,19 +24,12 @@ class PooledObjectPlacer:
     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__(
@@ -61,58 +54,79 @@ def __init__(
         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)
-
-        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
-
-            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"Placement pool 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 ones 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 candidates pass strict validation, the best-loss candidates 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)
 
@@ -126,23 +140,25 @@ def _solve_and_store(self, num_layouts: int) -> None:
             )
 
         if valid_results:
-            self._layouts.extend(valid_results)
-        else:
-            print("Warning: No candidates passed strict validation. Accepting best-loss layouts as fallback.")
-            self._layouts.extend(all_results)
+            return valid_results
 
-    # ------------------------------------------------------------------
-    # Heterogeneous (per-env pool) internals
-    # ------------------------------------------------------------------
+        print("Warning: No candidates passed strict validation. Accepting best-loss layouts as fallback.")
+        return all_results
 
-    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 _store_reusable_results(self, layouts: list[PlacementResult]) -> None:
+        """Distribute reusable layouts across env pools without dropping valid results."""
+        if not layouts:
+            return
+
+        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
@@ -151,17 +167,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)
@@ -175,142 +188,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)
-
-        remaining = len(self._layouts) - self._next_idx
-        if remaining < count:
-            self._solve_and_store(max(self._pool_size, count))
+        if count % self._num_envs != 0:
+            raise ValueError(f"count must be a multiple of num_envs ({self._num_envs}), got {count}")
 
-        remaining = len(self._layouts) - self._next_idx
-        if remaining < count:
-            raise RuntimeError(
-                f"Placement pool has {remaining} valid layouts but {count} were requested. "
-                "The solver is not producing enough valid placements."
-            )
-
-        start = self._next_idx
-        self._next_idx += count
-        return self._layouts[start : self._next_idx]
-
-    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
+        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
 
@@ -318,8 +292,7 @@ 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)
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 59bd9f933..fa2779e70 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 a684bfc4ab6f33f62eed42fe1de0158879eba703 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 14:57:04 -0700
Subject: [PATCH 18/24] 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 85e92b1ea..1ec9498e8 100644
--- a/isaaclab_arena/assets/object_base.py
+++ b/isaaclab_arena/assets/object_base.py
@@ -70,11 +70,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 d30126c1e..256905b92 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 = 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:
@@ -203,7 +203,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 fa2779e70..66ae2e757 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 72dd961ab62d8885f796ca3ca422aeedf60a48dd Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 15:23:12 -0700
Subject: [PATCH 19/24] 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 5ec80e15cb105c1c283da51a53fa993b5e942e88 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 17:11:21 -0700
Subject: [PATCH 20/24] 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 6cae73f7a93f5c95b7d1f6a57936ba8832fa23a9 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Mon, 11 May 2026 21:19:33 -0700
Subject: [PATCH 21/24] 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 01b97f7dce9479a925f568ce1e47ba3537159d7c Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 12 May 2026 13:50:48 -0700
Subject: [PATCH 22/24] 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 256905b92..2d6a3dfe2 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 maintains a pool of valid placement layouts.
 
-    Wraps :class:`ObjectPlacer` and solves 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 = 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 = 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]:
@@ -146,16 +151,21 @@ def _solve_reusable_layouts(self, num_layouts: int) -> list[PlacementResult]:
         return all_results
 
     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.
@@ -163,21 +173,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(
@@ -191,7 +201,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
@@ -210,36 +227,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.
@@ -248,41 +264,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))
@@ -290,9 +300,11 @@ 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())

From edbebd103a2631fe245147e18147b636e5c10110 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Tue, 12 May 2026 17:18:42 -0700
Subject: [PATCH 23/24] 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 | 25 +++++++
 5 files changed, 135 insertions(+), 27 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 7a53f0598..dce4ff49d 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 pooled placer 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 2d6a3dfe2..70e7ae69f 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:
@@ -247,20 +250,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}")
 
@@ -281,6 +290,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).
 
@@ -308,3 +348,8 @@ def remaining(self) -> int:
         minimum is what limits without-replacement capacity).
         """
         return min(self._available_per_env())
+
+    @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 66ae2e757..5a79631da 100644
--- a/isaaclab_arena/tests/test_placement_events.py
+++ b/isaaclab_arena/tests/test_placement_events.py
@@ -247,6 +247,31 @@ def test_solve_and_place_objects_handles_multiple_env_ids():
         )
 
 
+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 cf5599520cc013198f358e7aeba7b7a5b51dce56 Mon Sep 17 00:00:00 2001
From: zhx06 <zihaox@nvidia.com>
Date: Wed, 13 May 2026 09:13:05 -0700
Subject: [PATCH 24/24] 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 70e7ae69f..63d7d2956 100644
--- a/isaaclab_arena/relations/pooled_object_placer.py
+++ b/isaaclab_arena/relations/pooled_object_placer.py
@@ -158,8 +158,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
@@ -325,9 +325,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()