Add heterogeneous object placement support

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):
     """

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:
+        """Mirror ObjectBase.get_bounding_box_per_env for this test double."""
+        bbox = self.get_bounding_box()
+        return AxisAlignedBoundingBox(
+            min_point=bbox.min_point.expand(num_envs, 3),
+            max_point=bbox.max_point.expand(num_envs, 3),
+        )
+
     def get_world_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get bounding box in world coordinates (local bbox rotated and translated).
 

isaaclab_arena/assets/object_base.py

@@ -69,6 +69,27 @@ def get_world_bounding_box(self) -> AxisAlignedBoundingBox:
         """Get bounding box in world coordinates (local bbox rotated and translated)."""
         ...
 
+    def get_bounding_box_per_env(self, num_envs: int) -> AxisAlignedBoundingBox:
+        """Get local bounding boxes for each environment.
+
+        This default implementation is for objects with the same geometry in
+        every environment: it expands the single local bbox to ``(num_envs, 3)``.
+        ``RigidObjectSet`` overrides this to return the bbox for each env's
+        assigned variant.
+
+        Args:
+            num_envs: Number of environments.
+
+        Returns:
+            ``AxisAlignedBoundingBox`` with ``min_point`` / ``max_point`` of shape
+            ``(num_envs, 3)``.
+        """
+        bbox = self.get_bounding_box()
+        return AxisAlignedBoundingBox(
+            min_point=bbox.min_point.expand(num_envs, 3),
+            max_point=bbox.max_point.expand(num_envs, 3),
+        )
+
     def _get_initial_pose_as_pose(self) -> Pose | None:
         """Return a single ``Pose`` suitable for *init_state* and bounding-box calculations.
 

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
@@ -28,6 +30,7 @@ def __init__(
         prim_path: str | None = None,
         scale: tuple[float, float, float] = (1.0, 1.0, 1.0),
         random_choice: bool = False,
+        variant_indices_by_env: list[int] | None = None,
         initial_pose: Pose | None = None,
         **kwargs,
     ):
@@ -40,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. If False, variants are assigned by repeating
+                the member order across environments.
+            variant_indices_by_env: Optional fixed variant index for each environment.
             initial_pose: The initial pose of the object from this object set.
         """
         if not self._are_all_objects_type_rigid(objects):
@@ -63,10 +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 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:
@@ -90,27 +104,100 @@ 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.
+
+        Multi-variant sets use one fixed assignment for the lifetime of the
+        object set. When ``random_choice`` is True, each env independently
+        samples one variant once. Otherwise, assignments repeat the member
+        order across environments.
+
+        Args:
+            num_envs: Number of environments.
+
+        Returns:
+            List of length ``num_envs`` with indices into ``self.objects``.
+        """
+        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.
+
+        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.
         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:
+            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:
+        n = len(self.objects)
+        if any(idx < 0 or idx >= n for idx in variant_indices_by_env):
+            raise ValueError(
+                f"RigidObjectSet '{self.name}' variant indices must be in [0, {n}); got {variant_indices_by_env}."
+            )
+
+        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
         object_cfg = 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):
@@ -143,6 +230,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)

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:

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] | None = None,
+        env_cfg_callback: (
+            Callable[[IsaacLabArenaManagerBasedRLEnvCfg], IsaacLabArenaManagerBasedRLEnvCfg] | None
+        ) = None,
         rl_framework_entry_point: str | None = None,
         rl_policy_cfg: str | None = None,
     ):