improved bin packing example

ooctipus · ooctipus · commit e5b71fcca1cb · 2025-10-22T14:43:17.000-07:00
diff --git a/scripts/demos/bin_packing.py b/scripts/demos/bin_packing.py
@@ -8,7 +8,7 @@
 .. code-block:: bash
 
     # Usage
-    ./isaaclab.sh -p scripts/demos/bin_packing.py --num_envs 2048
+    ./isaaclab.sh -p scripts/demos/bin_packing.py --num_envs 32
 
 """
 
@@ -22,9 +22,7 @@
 from isaaclab.app import AppLauncher
 
 # add argparse arguments
-parser = argparse.ArgumentParser(
-    description="Demo on spawning different number of objects in multiple bin packing environments."
-)
+parser = argparse.ArgumentParser(description="Demo usage of RigidObjectCollection through bin packing example")
 parser.add_argument("--num_envs", type=int, default=16, help="Number of environments to spawn.")
 # append AppLauncher cli args
 AppLauncher.add_app_launcher_args(parser)
@@ -37,11 +35,12 @@
 
 """Rest everything follows."""
 
+import math
 import torch
 
 import isaaclab.sim as sim_utils
 import isaaclab.utils.math as math_utils
-from isaaclab.assets import AssetBaseCfg, RigidObject, RigidObjectCfg, RigidObjectCollection, RigidObjectCollectionCfg
+from isaaclab.assets import AssetBaseCfg, RigidObjectCfg, RigidObjectCollection, RigidObjectCollectionCfg
 from isaaclab.scene import InteractiveScene, InteractiveSceneCfg
 from isaaclab.sim import SimulationContext
 from isaaclab.utils import Timer, configclass
@@ -51,42 +50,34 @@
 # Scene Configuration
 ##
 
+MAX_NUM_OBJECTS = 24
+MAX_OBJECTS_PER_BIN = 24
+MIN_OBJECTS_PER_BIN = 1
+NUM_OBJECTS_PER_LAYER = 4
+CACHE_HEIGHT = 2.5
 POSE_RANGE = {"roll": (-3.14, 3.14), "pitch": (-3.14, 3.14), "yaw": (-3.14, 3.14)}
-VELOCITY_RANGE = {"roll": (-0.76, 0.76), "pitch": (-0.76, 0.76), "yaw": (-0.76, 0.76)}
+VELOCITY_RANGE = {"roll": (-0.2, 1.0), "pitch": (-0.2, 1.0), "yaw": (-0.2, 1.0)}
 
+# Object layout configuration
 
-RANDOM_YCB_RIGID_OBJECT_CFG = RigidObjectCfg(
-    prim_path="/World/envs/env_.*/Object",
-    spawn=sim_utils.MultiAssetSpawnerCfg(
-        assets_cfg=[
-            sim_utils.UsdFileCfg(
-                usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/004_sugar_box.usd",
-                rigid_props=sim_utils.RigidBodyPropertiesCfg(solver_position_iteration_count=4),
-            ),
-            sim_utils.UsdFileCfg(
-                usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/003_cracker_box.usd",
-                rigid_props=sim_utils.RigidBodyPropertiesCfg(solver_position_iteration_count=4),
-            ),
-            sim_utils.UsdFileCfg(
-                usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/005_tomato_soup_can.usd",
-                rigid_props=sim_utils.RigidBodyPropertiesCfg(solver_position_iteration_count=4),
-            ),
-            sim_utils.UsdFileCfg(
-                usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/006_mustard_bottle.usd",
-                rigid_props=sim_utils.RigidBodyPropertiesCfg(solver_position_iteration_count=4),
-            ),
-            # note: the placeholder, this allows the effect of having less objects in some env ids
-            sim_utils.SphereCfg(
-                radius=0.1,
-                rigid_props=sim_utils.RigidBodyPropertiesCfg(disable_gravity=True),
-                collision_props=sim_utils.CollisionPropertiesCfg(collision_enabled=False),
-                visible=False
-            ),
-        ],
-        random_choice=True,
+GROCERIES = {
+    "OBJECT_A" : sim_utils.UsdFileCfg(
+        usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/004_sugar_box.usd",
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(solver_position_iteration_count=4),
+    ),
+    "OBJECT_B" : sim_utils.UsdFileCfg(
+        usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/003_cracker_box.usd",
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(solver_position_iteration_count=4),
+    ),
+    "OBJECT_C" : sim_utils.UsdFileCfg(
+        usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/005_tomato_soup_can.usd",
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(solver_position_iteration_count=4),
     ),
-    init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 2.0)),
-)
+    "OBJECT_D" : sim_utils.UsdFileCfg(
+        usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/006_mustard_bottle.usd",
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(solver_position_iteration_count=4),
+    )
+}
 
 
 @configclass
@@ -115,82 +106,103 @@ class MultiObjectSceneCfg(InteractiveSceneCfg):
         init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.15)),
     )
 
-    # object collection
-    object_collection: RigidObjectCollectionCfg = RigidObjectCollectionCfg(
+    groceries: RigidObjectCollectionCfg = RigidObjectCollectionCfg(
         rigid_objects={
-            "Object_A_Layer1": RANDOM_YCB_RIGID_OBJECT_CFG.replace(
-                prim_path="/World/envs/env_.*/Object_A_Layer1",
-                init_state=RigidObjectCfg.InitialStateCfg(pos=(-0.035, -0.06, 0.2)),
-            ),
-            "Object_B_Layer1": RANDOM_YCB_RIGID_OBJECT_CFG.replace(
-                prim_path="/World/envs/env_.*/Object_B_Layer1",
-                init_state=RigidObjectCfg.InitialStateCfg(pos=(-0.035, 0.06, 0.2)),
-            ),
-            "Object_C_Layer1": RANDOM_YCB_RIGID_OBJECT_CFG.replace(
-                prim_path="/World/envs/env_.*/Object_C_Layer1",
-                init_state=RigidObjectCfg.InitialStateCfg(pos=(0.035, 0.06, 0.2)),
-            ),
-            "Object_D_Layer1": RANDOM_YCB_RIGID_OBJECT_CFG.replace(
-                prim_path="/World/envs/env_.*/Object_D_Layer1",
-                init_state=RigidObjectCfg.InitialStateCfg(pos=(0.035, -0.06, 0.2)),
-            ),
-            "Object_A_Layer2": RANDOM_YCB_RIGID_OBJECT_CFG.replace(
-                prim_path="/World/envs/env_.*/Object_A_Layer2",
-                init_state=RigidObjectCfg.InitialStateCfg(pos=(-0.035, -0.06, 0.4)),
-            ),
-            "Object_B_Layer2": RANDOM_YCB_RIGID_OBJECT_CFG.replace(
-                prim_path="/World/envs/env_.*/Object_B_Layer2",
-                init_state=RigidObjectCfg.InitialStateCfg(pos=(-0.035, 0.06, 0.4)),
-            ),
-            "Object_C_Layer2": RANDOM_YCB_RIGID_OBJECT_CFG.replace(
-                prim_path="/World/envs/env_.*/Object_C_Layer2",
-                init_state=RigidObjectCfg.InitialStateCfg(pos=(0.035, 0.06, 0.4)),
-            ),
-            "Object_D_Layer2": RANDOM_YCB_RIGID_OBJECT_CFG.replace(
-                prim_path="/World/envs/env_.*/Object_D_Layer2",
-                init_state=RigidObjectCfg.InitialStateCfg(pos=(0.035, -0.06, 0.4)),
-            ),
+            f"Object_{label}_Layer{layer}": RigidObjectCfg(
+                prim_path=f"/World/envs/env_.*/Object_{label}_Layer{layer}",
+                init_state=RigidObjectCfg.InitialStateCfg(
+                    pos=(x, y, 0.2 + (layer) * 0.2)
+                ),
+                spawn=GROCERIES.get(f"OBJECT_{label}")
+            )
+            for layer in range(MAX_NUM_OBJECTS // NUM_OBJECTS_PER_LAYER)
+            for label, (x, y) in zip(["A", "B", "C", "D"], [(-0.035, -0.1), (-0.035, 0.1), (0.035, 0.1), (0.035, -0.1)])
         }
     )
 
-def reset_object_collections(scene: InteractiveScene, view_ids: torch.Tensor):
-    if len(view_ids) == 0:
-        return
-    rigid_object_collection: RigidObjectCollection = scene["object_collection"]
-    default_state_w = rigid_object_collection.data.default_object_state.clone()
-    default_state_w[..., :3] = default_state_w[..., :3] + scene.env_origins.unsqueeze(1)
-    default_state_w_view = rigid_object_collection.reshape_data_to_view(default_state_w)[view_ids]
-    range_list = [POSE_RANGE.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
-    ranges = torch.tensor(range_list, device=scene.device)
-    samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(view_ids), 6), device=scene.device)
 
+def reset_object_collections(
+    scene: InteractiveScene, collection_name: str, view_states: torch.Tensor, view_ids: torch.Tensor, randomize: bool = False
+):
+    rigid_object_collection: RigidObjectCollection = scene[collection_name]
+    sel_view_states = view_states[view_ids]
+    positions = sel_view_states[:, :3]
+    orientations = sel_view_states[:, 3:7]
     # poses
-    positions = default_state_w_view[:, :3] + samples[..., 0:3]
-    orientations_delta = math_utils.quat_from_euler_xyz(samples[..., 3], samples[..., 4], samples[..., 5])
-    orientations = math_utils.quat_mul(default_state_w_view[:, 3:7], orientations_delta)
-    # velocities
-    range_list = [VELOCITY_RANGE.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
-    ranges = torch.tensor(range_list, device=scene.device)
-    samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(view_ids), 6), device=scene.device)
-
-    new_velocities = default_state_w_view[:, 7:13] + samples
-    new_poses = torch.concat((positions, orientations), dim=-1)
-
-    new_poses[..., 3:] = math_utils.convert_quat(new_poses[..., 3:], to="xyzw")
+    if randomize:
+        range_list = [POSE_RANGE.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
+        ranges = torch.tensor(range_list, device=scene.device)
+        samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(view_ids), 6), device=scene.device)
+        positions += samples[..., 0:3]
 
-    num_objects = rigid_object_collection.num_instances * rigid_object_collection.num_objects
+        orientations_delta = math_utils.quat_from_euler_xyz(samples[..., 3], samples[..., 4], samples[..., 5])
+        orientations = math_utils.convert_quat(orientations, to="wxyz")
+        orientations = math_utils.quat_mul(orientations, orientations_delta)
+        orientations = math_utils.convert_quat(orientations, to="xyzw")
 
-    if len(view_ids) != num_objects:
-        poses = torch.zeros((num_objects, 7), device=scene.device)
-        poses[view_ids, :] = new_poses
-        velocities = torch.zeros((num_objects, 6), device=scene.device)
-        velocities[view_ids, :] = new_velocities
+    # velocities
+    new_velocities = sel_view_states[:, 7:13]
+    if randomize:
+        range_list = [VELOCITY_RANGE.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
+        ranges = torch.tensor(range_list, device=scene.device)
+        samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(view_ids), 6), device=scene.device)
+        new_velocities += samples
     else:
-        poses = new_poses
-        velocities = new_velocities
-
-    rigid_object_collection.root_physx_view.set_transforms(poses, indices=view_ids.view(-1, 1))
-    rigid_object_collection.root_physx_view.set_velocities(velocities, indices=view_ids.view(-1, 1))
+        new_velocities[:] = 0.0
+
+    view_states[view_ids, :7] = torch.concat((positions, orientations), dim=-1)
+    view_states[view_ids, 7:] = new_velocities
+
+    rigid_object_collection.root_physx_view.set_transforms(view_states[:, :7], indices=view_ids)
+    rigid_object_collection.root_physx_view.set_velocities(view_states[:, 7:], indices=view_ids)
+
+
+def build_grocery_defaults(
+    num_envs: int,
+    device: str = "cpu",
+    z_spacing: float = 0.1,
+    cache_spacing: float = 0.25,
+    num_objects_per_layer: int = NUM_OBJECTS_PER_LAYER,
+    max_objects_per_bin: int = MAX_OBJECTS_PER_BIN,
+    max_num_objects: int = MAX_NUM_OBJECTS,
+    cache_height: float = CACHE_HEIGHT
+):
+
+    # The bin has a size of 0.2 x 0.3 x 0.15 m
+    bin_x_dim, bin_y_dim, bin_z_dim = 0.2, 0.3, 0.15
+    # First, we calculate the number of layers and objects per layer
+    num_layers = math.ceil(max_objects_per_bin / num_objects_per_layer)
+    num_x_objects = math.ceil(math.sqrt(num_objects_per_layer))
+    num_y_objects = math.ceil(num_objects_per_layer / num_x_objects)
+    total_objects = num_x_objects * num_y_objects * num_layers
+    # Then, we create a 3D grid that allows for IxJxN objects to be placed ontop of the bin.
+    x = torch.linspace(-bin_x_dim * (2 / 6), bin_x_dim * (2 / 6), num_x_objects, device=device)
+    y = torch.linspace(-bin_y_dim * (2 / 6), bin_y_dim * (2 / 6), num_y_objects, device=device)
+    z = torch.linspace(0, z_spacing * (num_layers - 1), num_layers, device=device) + bin_z_dim * 2
+    grid_z, grid_y, grid_x = torch.meshgrid(z, y, x, indexing="ij")  # Note Z first, allowing to naturally stack the layers.
+    # Using this grid plus a reference quaternion, we can create the poses for the groceries to be spawned above the bin.
+    ref_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0]], device=device).repeat(total_objects, 1)
+    positions = torch.stack((grid_x.flatten(), grid_y.flatten(), grid_z.flatten()), dim=-1)
+    poses = torch.cat((positions, ref_quat), dim=-1)
+    # Duplicate across environments, cap at max_num_objects
+    active_spawn_poses = poses.unsqueeze(0).repeat(num_envs, 1, 1)[:, :max_num_objects, :]
+
+    # We'll also create a buffer for the cached groceries. They'll be spawned below the bin so they can't be seen.
+    num_x_objects = math.ceil(math.sqrt(max_num_objects))
+    num_y_objects = math.ceil(max_num_objects / num_x_objects)
+    # We create a XY grid only and fix the Z height for the cache.
+    x = cache_spacing * torch.arange(num_x_objects, device=device)
+    y = cache_spacing * torch.arange(num_y_objects, device=device)
+    grid_y, grid_x = torch.meshgrid(y, x, indexing="ij")
+    grid_z = cache_height * torch.ones_like(grid_x)
+    # We can then create the poses for the cached groceries.
+    ref_quat = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=device).repeat(num_x_objects * num_y_objects, 1)
+    positions = torch.stack((grid_x.flatten(), grid_y.flatten(), grid_z.flatten()), dim=-1)
+    poses = torch.cat((positions, ref_quat), dim=-1)
+    # Duplicate across environments, cap at max_num_objects
+    cached_spawn_poses = poses.unsqueeze(0).repeat(num_envs, 1, 1)[:, :max_num_objects, :]
+
+    return active_spawn_poses, cached_spawn_poses
 
 
 ##
@@ -202,24 +214,38 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene):
     """Runs the simulation loop."""
     # Extract scene entities
     # note: we only do this here for readability.
-    rigid_object: RigidObject = scene["object"]
-    rigid_object_collection: RigidObjectCollection = scene["object_collection"]
+    rigid_object_collection: RigidObjectCollection = scene["groceries"]
     view_indices = torch.arange(scene.num_envs * rigid_object_collection.num_objects, device=scene.device)
+    default_state_w = rigid_object_collection.data.default_object_state.clone()
+    default_state_w[..., :3] = default_state_w[..., :3] + scene.env_origins.unsqueeze(1)
     # Define simulation stepping
     sim_dt = sim.get_physics_dt()
     count = 0
+
+    active_spawn_poses, cached_spawn_poses = build_grocery_defaults(scene.num_envs, scene.device, max_num_objects=rigid_object_collection.num_objects)
+    active_spawn_poses[..., :3] += scene.env_origins.view(-1, 1, 3)
+    cached_spawn_poses[..., :3] += scene.env_origins.view(-1, 1, 3)
+    active_spawn_poses = rigid_object_collection.reshape_data_to_view(active_spawn_poses)
+    cached_spawn_poses = rigid_object_collection.reshape_data_to_view(cached_spawn_poses)
+    spawning_state_w_view_xyzw = rigid_object_collection.reshape_data_to_view(default_state_w).clone()
+
+    groceries_mask_helper = torch.arange(rigid_object_collection.num_objects * scene.num_envs, device=scene.device) % rigid_object_collection.num_objects
     # Simulation loop
     while simulation_app.is_running():
         # Reset
         if count % 250 == 0:
             # reset counter
             count = 0
-            # reset the scene entities
-            # object
-            root_state = rigid_object.data.default_root_state.clone()
-            root_state[:, :3] += scene.env_origins
-            # object collection
-            reset_object_collections(scene, view_indices)
+            num_active_groceries = torch.randint(MIN_OBJECTS_PER_BIN, rigid_object_collection.num_objects, (scene.num_envs, 1), device=scene.device)
+            groceries_mask = (groceries_mask_helper.view(scene.num_envs, -1) < num_active_groceries).view(-1, 1)
+            spawning_state_w_view_xyzw[:, :7] = cached_spawn_poses * (~groceries_mask) + active_spawn_poses * groceries_mask
+            # Retrieve positions
+            reset_object_collections(scene, "groceries", spawning_state_w_view_xyzw, view_indices[~groceries_mask.view(-1)])
+            reset_object_collections(scene, "groceries", spawning_state_w_view_xyzw, view_indices[groceries_mask.view(-1)], randomize=True)
+            # reset_object_collections(scene, "groceries", spawning_state_w_view_xyzw, view_indices)
+            random_masses = torch.rand(rigid_object_collection.num_instances * rigid_object_collection.num_objects, device=scene.device) * 0.2 + 0.2
+            rigid_object_collection.root_physx_view.set_masses(random_masses.cpu(), view_indices.cpu())
+            rigid_object_collection.root_physx_view.set_disable_gravities((~groceries_mask).cpu(), indices=view_indices.cpu())
             scene.reset()
             print("[INFO]: Resetting scene state...")
 
@@ -231,7 +257,9 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene):
         object_pos_b_view = rigid_object_collection.reshape_data_to_view(object_pos_b)
         inbound_mask = (-0.2 < object_pos_b_view[:, 0]) & (object_pos_b_view[:, 0] < 0.2)
         inbound_mask &= (-0.3 < object_pos_b_view[:, 1]) & (object_pos_b_view[:, 1] < 0.3)
-        reset_object_collections(scene, view_indices[~inbound_mask])
+        out_bound_indices = view_indices[~inbound_mask]
+        if len(out_bound_indices) > 0:
+            reset_object_collections(scene, "groceries", spawning_state_w_view_xyzw, view_indices[out_bound_indices])
         # Increment counter
         count += 1
         # Update buffers
