Added randomizing rigid object material to tensor demo

Author: ooctipus

scripts/demos/amz_tensor_demo.py

@@ -67,6 +67,11 @@
 
 CACHE_HEIGHT = 2.0
 
+STATIC_FRICTION_RANGE = (0.2, 2.0)
+DYNAMIC_FRICTION_RANGE = (0.2, 2.0)
+RESTITUTION_RANGE = (0.0, 0.0)
+NUM_BUCKETS = 25
+
 def setup_object_cfg(max_num_objects: int, spacing: float = 0.5) -> RigidObjectCollectionCfg:
     object_cfg_list = []
     grid_size = math.ceil(math.sqrt(max_num_objects))
@@ -187,6 +192,11 @@ def __init__(self,
         self.current_grocery_poses = torch.zeros((self.num_envs, MAX_NUM_OBJECTS, 7), device=self.sim.device)
         self.current_grocery_velocities = torch.zeros((self.num_envs, MAX_NUM_OBJECTS, 6), device=self.sim.device)
 
+        range_list = [STATIC_FRICTION_RANGE, DYNAMIC_FRICTION_RANGE, RESTITUTION_RANGE]
+        ranges = torch.tensor(range_list, device="cpu")
+        self.material_buckets = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (NUM_BUCKETS, 3), device="cpu")
+
+
         self.ALL_OBJ_INDICES = torch.arange(MAX_NUM_OBJECTS, dtype=torch.long, device=self.sim.device)
         self.ALL_ENV_INDICES = torch.arange(self.num_envs, dtype=torch.long, device=self.sim.device)
 
@@ -287,6 +297,21 @@ def set_mass_tensor(self, mass, env_ids: list[int] | None = None):
             reworked_indices = (self.ALL_OBJ_INDICES.unsqueeze(1) * num_envs + env_ids).flatten()
             self.groceries_view.set_masses(mass.flatten().cpu(), indices=reworked_indices.cpu())
 
+    def set_rigid_body_material_tensor(self, materials: torch.Tensor, env_ids: list[int] | None = None):
+        with Timer("[INFO] Time to set the rigid_object_material"):
+            # Set the rigid body material of the groceries in the view
+            # Doc here: https://docs.omniverse.nvidia.com/kit/docs/omni_physics/latest/extensions/runtime/source/omni.physics.tensors/docs/api/python.html#omni.physics.tensors.impl.api.ArticulationView.set_material_properties
+
+            # Note: We need to convert to CPU as this property can only be set on CPU,
+            # we need to rework the indices to match the shape of the view.
+            if env_ids is None:
+                num_envs = self.num_envs
+            else:
+                num_envs = len(env_ids)
+            
+            reworked_indices = (self.ALL_OBJ_INDICES.unsqueeze(1) * num_envs + env_ids).flatten()
+            self.groceries_view.set_material_properties(materials, indices=reworked_indices.cpu())
+
     def set_disable_gravity_tensor(self, disable_gravity: torch.Tensor, env_ids: list[int] | None = None):
         with Timer("[INFO] Time to set the disable gravity"):
             # Disables the gravity for some of the inactive groceries.
@@ -411,9 +436,19 @@ def reset(self, env_ids: list[int] | None = None):
 
         # Randomize the mass of the groceries
         new_masses = torch.rand((len(env_ids), MAX_NUM_OBJECTS), device=self.sim.device) * 0.2 + 0.2
+
+        # Randomize the grocery materials
+        total_num_shapes = self.groceries_view.max_shapes
+        bucket_ids = torch.randint(0, len(self.material_buckets), (len(env_ids), total_num_shapes), device="cpu")
+        material_samples = self.material_buckets[bucket_ids]
+
+        # retrieve material buffer from the physics simulation
+        materials = self.groceries_view.get_material_properties()
+        materials[env_ids] = material_samples[:]
+
         # Apply the new masses to the simulation
         self.set_mass_tensor(new_masses, env_ids)
-
+        self.set_rigid_body_material_tensor(materials, env_ids)
         # Disable the gravity for the cached groceries
         self.set_disable_gravity_tensor(~self.groceries_mask[env_ids], env_ids)
 

scripts/demos/bin_packing.py

@@ -153,7 +153,6 @@ class MultiObjectSceneCfg(InteractiveSceneCfg):
         }
     )
 
-
 def reset_object_collections(scene: InteractiveScene, view_ids: torch.Tensor):
     if len(view_ids) == 0:
         return
@@ -165,6 +164,7 @@ def reset_object_collections(scene: InteractiveScene, view_ids: torch.Tensor):
     ranges = torch.tensor(range_list, device=scene.device)
     samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(view_ids), 6), device=scene.device)
 
+    # 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)
@@ -173,13 +173,24 @@ def reset_object_collections(scene: InteractiveScene, view_ids: torch.Tensor):
     ranges = torch.tensor(range_list, device=scene.device)
     samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(view_ids), 6), device=scene.device)
 
-    velocities = default_state_w_view[:, 7:13] + samples
+    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")
-    rigid_object_collection.root_physx_view.set_transforms(new_poses, indices=view_ids.view(-1, 1))
-    rigid_object_collection.root_physx_view.set_velocities(velocities, indices=view_ids.view(-1, 1))
 
+    num_objects = rigid_object_collection.num_instances * rigid_object_collection.num_objects
+
+    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
+    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))
 
 
 ##
@@ -207,8 +218,6 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene):
             # object
             root_state = rigid_object.data.default_root_state.clone()
             root_state[:, :3] += scene.env_origins
-            rigid_object.write_root_pose_to_sim(root_state[:, :7])
-            rigid_object.write_root_velocity_to_sim(root_state[:, 7:])
             # object collection
             reset_object_collections(scene, view_indices)
             scene.reset()
@@ -220,8 +229,8 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene):
         sim.step()
         object_pos_b = rigid_object_collection.data.object_pos_w - scene.env_origins.unsqueeze(1)
         object_pos_b_view = rigid_object_collection.reshape_data_to_view(object_pos_b)
-        inbound_mask = (-1.0 < object_pos_b_view[:, 0]) & (object_pos_b_view[:, 0] < 1.0)
-        inbound_mask &= (-1.0 < object_pos_b_view[:, 1]) & (object_pos_b_view[:, 1] < 1.0)
+        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])
         # Increment counter
         count += 1