Fix/playback scale

omnigibson/envs/data_wrapper.py

@@ -293,6 +293,9 @@ def __init__(
         # Denotes the maximum serialized state size for the current episode
         self.max_state_size = 0
 
+        # Dict capturing serialized per-episode initial information (e.g.: scales / visibilities) about every object
+        self.init_metadata = dict()
+
         # Maps episode step ID to dictionary of systems and objects that should be added / removed to the simulator at
         # the given simulator step. See add_transition_info() for more info
         self.current_transitions = dict()
@@ -390,6 +393,19 @@ def reset(self):
         # Update max state size
         self.max_state_size = max(self.max_state_size, len(state))
 
+        # Also store initial metadata (scale, visibility) not recorded in serialized state
+        # This is simply serialized
+        scales = th.zeros(self.scene.n_objects, 3)
+        visibilities = th.zeros(self.scene.n_objects, dtype=th.bool)
+        for i, obj in enumerate(self.scene.objects):
+            scales[i] = obj.scale
+            visibilities[i] = obj.visible
+
+        self.init_metadata = {
+            "scales": scales,
+            "visibilities": visibilities,
+        }
+
         return init_obs, init_info
 
     def _parse_step_data(self, action, obs, reward, terminated, truncated, info):
@@ -422,6 +438,10 @@ def process_traj_to_hdf5(self, traj_data, traj_grp_name, nested_keys=("obs",)):
         # Add in transition info
         self.add_metadata(group=traj_grp, name="transitions", data=self.current_transitions)
 
+        # Add initial metadata information
+        for name, data in self.init_metadata.items():
+            traj_grp.create_dataset(name, data=data)
+
         return traj_grp
 
     def flush_current_traj(self):
@@ -634,6 +654,8 @@ def playback_episode(self, episode_id, record_data=True, video_writer=None, vide
         # Grab episode data
         transitions = json.loads(traj_grp.attrs["transitions"])
         traj_grp = h5py_group_to_torch(traj_grp)
+        scales = traj_grp["scales"]
+        visibilities = traj_grp["visibilities"]
         action = traj_grp["action"]
         state = traj_grp["state"]
         state_size = traj_grp["state_size"]
@@ -643,6 +665,13 @@ def playback_episode(self, episode_id, record_data=True, video_writer=None, vide
 
         # Reset environment
         og.sim.restore(scene_files=[self.scene_file])
+
+        # Reset scales and visibilities
+        with og.sim.stopped():
+            assert len(scales) == self.scene.n_objects and len(visibilities) == self.scene.n_objects
+            for scale, visible, obj in zip(scales, visibilities, self.scene.objects):
+                obj.scale = scale
+                obj.visible = bool(visible)
         self.reset()
 
         # Restore to initial state

omnigibson/objects/controllable_object.py

@@ -42,6 +42,7 @@ def __init__(
         visual_only=False,
         self_collisions=False,
         prim_type=PrimType.RIGID,
+        link_physics_materials=None,
         load_config=None,
         control_freq=None,
         controller_config=None,
@@ -63,6 +64,10 @@ def __init__(
             visual_only (bool): Whether this object should be visual only (and not collide with any other objects)
             self_collisions (bool): Whether to enable self collisions for this object
             prim_type (PrimType): Which type of prim the object is, Valid options are: {PrimType.RIGID, PrimType.CLOTH}
+            link_physics_materials (None or dict): If specified, dictionary mapping link name to kwargs used to generate
+                a specific physical material for that link's collision meshes, where the kwargs are arguments directly
+                passed into the omni.isaac.core.materials.PhysicsMaterial constructor, e.g.: "static_friction",
+                "dynamic_friction", and "restitution"
             load_config (None or dict): If specified, should contain keyword-mapped values that are relevant for
                 loading this prim at runtime.
             control_freq (float): control frequency (in Hz) at which to control the object. If set to be None,
@@ -125,6 +130,7 @@ def __init__(
             visual_only=visual_only,
             self_collisions=self_collisions,
             prim_type=prim_type,
+            link_physics_materials=link_physics_materials,
             load_config=load_config,
             **kwargs,
         )

omnigibson/objects/dataset_object.py

@@ -50,6 +50,7 @@ def __init__(
         kinematic_only=None,
         self_collisions=False,
         prim_type=PrimType.RIGID,
+        link_physics_materials=None,
         load_config=None,
         abilities=None,
         include_default_states=True,
@@ -82,6 +83,10 @@ def __init__(
                 are satisfied (see object_base.py post_load function), else False.
             self_collisions (bool): Whether to enable self collisions for this object
             prim_type (PrimType): Which type of prim the object is, Valid options are: {PrimType.RIGID, PrimType.CLOTH}
+            link_physics_materials (None or dict): If specified, dictionary mapping link name to kwargs used to generate
+                a specific physical material for that link's collision meshes, where the kwargs are arguments directly
+                passed into the omni.isaac.core.materials.PhysicsMaterial constructor, e.g.: "static_friction",
+                "dynamic_friction", and "restitution"
             load_config (None or dict): If specified, should contain keyword-mapped values that are relevant for
                 loading this prim at runtime.
             abilities (None or dict): If specified, manually adds specific object states to this object. It should be
@@ -146,6 +151,7 @@ def __init__(
             self_collisions=self_collisions,
             prim_type=prim_type,
             include_default_states=include_default_states,
+            link_physics_materials=link_physics_materials,
             load_config=load_config,
             abilities=abilities,
             **kwargs,

omnigibson/objects/light_object.py

@@ -35,6 +35,7 @@ def __init__(
         category="light",
         scale=None,
         fixed_base=False,
+        link_physics_materials=None,
         load_config=None,
         abilities=None,
         include_default_states=True,
@@ -52,6 +53,10 @@ def __init__(
                 for this object. A single number corresponds to uniform scaling along the x,y,z axes, whereas a
                 3-array specifies per-axis scaling.
             fixed_base (bool): whether to fix the base of this object or not
+            link_physics_materials (None or dict): If specified, dictionary mapping link name to kwargs used to generate
+                a specific physical material for that link's collision meshes, where the kwargs are arguments directly
+                passed into the omni.isaac.core.materials.PhysicsMaterial constructor, e.g.: "static_friction",
+                "dynamic_friction", and "restitution"
             load_config (None or dict): If specified, should contain keyword-mapped values that are relevant for
                 loading this prim at runtime.
             abilities (None or dict): If specified, manually adds specific object states to this object. It should be
@@ -88,6 +93,7 @@ def __init__(
             self_collisions=False,
             prim_type=PrimType.RIGID,
             include_default_states=include_default_states,
+            link_physics_materials=link_physics_materials,
             load_config=load_config,
             abilities=abilities,
             **kwargs,

omnigibson/objects/object_base.py

@@ -47,6 +47,7 @@ def __init__(
         kinematic_only=None,
         self_collisions=False,
         prim_type=PrimType.RIGID,
+        link_physics_materials=None,
         load_config=None,
         **kwargs,
     ):
@@ -66,6 +67,10 @@ def __init__(
                 are satisfied (see object_base.py post_load function), else False.
             self_collisions (bool): Whether to enable self collisions for this object
             prim_type (PrimType): Which type of prim the object is, Valid options are: {PrimType.RIGID, PrimType.CLOTH}
+            link_physics_materials (None or dict): If specified, dictionary mapping link name to kwargs used to generate
+                a specific physical material for that link's collision meshes, where the kwargs are arguments directly
+                passed into the omni.isaac.core.materials.PhysicsMaterial constructor, e.g.: "static_friction",
+                "dynamic_friction", and "restitution"
             load_config (None or dict): If specified, should contain keyword-mapped values that are relevant for
                 loading this prim at runtime.
             kwargs (dict): Additional keyword arguments that are used for other super() calls from subclasses, allowing
@@ -80,6 +85,7 @@ def __init__(
         self._uuid = get_uuid(name, deterministic=True)
         self.category = category
         self.fixed_base = fixed_base
+        self._link_physics_materials = dict() if link_physics_materials is None else link_physics_materials
 
         # Values to be created at runtime
         self._highlight_cached_values = None
@@ -212,6 +218,20 @@ def _post_load(self):
             self.solver_position_iteration_count = m.DEFAULT_SOLVER_POSITION_ITERATIONS
             self.solver_velocity_iteration_count = m.DEFAULT_SOLVER_VELOCITY_ITERATIONS
 
+        # Add link materials if specified
+        if self._link_physics_materials is not None:
+            for link_name, material_info in self._link_physics_materials.items():
+                # We will permute the link materials dict in place to now point to the created material
+                mat_name = f"{link_name}_physics_mat"
+                physics_mat = lazy.omni.isaac.core.materials.PhysicsMaterial(
+                    prim_path=f"{self.prim_path}/Looks/{mat_name}",
+                    name=mat_name,
+                    **material_info,
+                )
+                for msh in self.links[link_name].collision_meshes.values():
+                    msh.apply_physics_material(physics_mat)
+                self._link_physics_materials[link_name] = physics_mat
+
         # Add semantics
         lazy.omni.isaac.core.utils.semantics.add_update_semantics(
             prim=self._prim,

omnigibson/objects/primitive_object.py

@@ -38,6 +38,7 @@ def __init__(
         kinematic_only=None,
         self_collisions=False,
         prim_type=PrimType.RIGID,
+        link_physics_materials=None,
         load_config=None,
         abilities=None,
         include_default_states=True,
@@ -65,6 +66,10 @@ def __init__(
                 are satisfied (see object_base.py post_load function), else False.
             self_collisions (bool): Whether to enable self collisions for this object
             prim_type (PrimType): Which type of prim the object is, Valid options are: {PrimType.RIGID, PrimType.CLOTH}
+            link_physics_materials (None or dict): If specified, dictionary mapping link name to kwargs used to generate
+                a specific physical material for that link's collision meshes, where the kwargs are arguments directly
+                passed into the omni.isaac.core.materials.PhysicsMaterial constructor, e.g.: "static_friction",
+                "dynamic_friction", and "restitution"
             load_config (None or dict): If specified, should contain keyword-mapped values that are relevant for
                 loading this prim at runtime.
             abilities (None or dict): If specified, manually adds specific object states to this object. It should be
@@ -109,6 +114,7 @@ def __init__(
             self_collisions=self_collisions,
             prim_type=prim_type,
             include_default_states=include_default_states,
+            link_physics_materials=link_physics_materials,
             load_config=load_config,
             abilities=abilities,
             **kwargs,

omnigibson/objects/stateful_object.py

@@ -75,6 +75,7 @@ def __init__(
         kinematic_only=None,
         self_collisions=False,
         prim_type=PrimType.RIGID,
+        link_physics_materials=None,
         load_config=None,
         abilities=None,
         include_default_states=True,
@@ -96,6 +97,10 @@ def __init__(
                 are satisfied (see object_base.py post_load function), else False.
             self_collisions (bool): Whether to enable self collisions for this object
             prim_type (PrimType): Which type of prim the object is, Valid options are: {PrimType.RIGID, PrimType.CLOTH}
+            link_physics_materials (None or dict): If specified, dictionary mapping link name to kwargs used to generate
+                a specific physical material for that link's collision meshes, where the kwargs are arguments directly
+                passed into the omni.isaac.core.materials.PhysicsMaterial constructor, e.g.: "static_friction",
+                "dynamic_friction", and "restitution"
             load_config (None or dict): If specified, should contain keyword-mapped values that are relevant for
                 loading this prim at runtime.
             abilities (None or dict): If specified, manually adds specific object states to this object. It should be
@@ -133,6 +138,7 @@ def __init__(
             kinematic_only=kinematic_only,
             self_collisions=self_collisions,
             prim_type=prim_type,
+            link_physics_materials=link_physics_materials,
             load_config=load_config,
             **kwargs,
         )

omnigibson/objects/usd_object.py

@@ -28,6 +28,7 @@ def __init__(
         kinematic_only=None,
         self_collisions=False,
         prim_type=PrimType.RIGID,
+        link_physics_materials=None,
         load_config=None,
         abilities=None,
         include_default_states=True,
@@ -51,6 +52,10 @@ def __init__(
                 are satisfied (see object_base.py post_load function), else False.
             self_collisions (bool): Whether to enable self collisions for this object
             prim_type (PrimType): Which type of prim the object is, Valid options are: {PrimType.RIGID, PrimType.CLOTH}
+            link_physics_materials (None or dict): If specified, dictionary mapping link name to kwargs used to generate
+                a specific physical material for that link's collision meshes, where the kwargs are arguments directly
+                passed into the omni.isaac.core.materials.PhysicsMaterial constructor, e.g.: "static_friction",
+                "dynamic_friction", and "restitution"
             load_config (None or dict): If specified, should contain keyword-mapped values that are relevant for
                 loading this prim at runtime.
             abilities (None or dict): If specified, manually adds specific object states to this object. It should be
@@ -76,6 +81,7 @@ def __init__(
             self_collisions=self_collisions,
             prim_type=prim_type,
             include_default_states=include_default_states,
+            link_physics_materials=link_physics_materials,
             load_config=load_config,
             abilities=abilities,
             **kwargs,

omnigibson/robots/a1.py

@@ -21,6 +21,7 @@ def __init__(
         visible=True,
         visual_only=False,
         self_collisions=True,
+        link_physics_materials=None,
         load_config=None,
         fixed_base=True,
         # Unique to USDObject hierarchy
@@ -39,6 +40,8 @@ def __init__(
         sensor_config=None,
         # Unique to ManipulationRobot
         grasping_mode="physical",
+        finger_static_friction=None,
+        finger_dynamic_friction=None,
         end_effector="inspire",
         **kwargs,
     ):
@@ -52,6 +55,10 @@ def __init__(
             visible (bool): whether to render this object or not in the stage
             visual_only (bool): Whether this object should be visual only (and not collide with any other objects)
             self_collisions (bool): Whether to enable self collisions for this object
+            link_physics_materials (None or dict): If specified, dictionary mapping link name to kwargs used to generate
+                a specific physical material for that link's collision meshes, where the kwargs are arguments directly
+                passed into the omni.isaac.core.materials.PhysicsMaterial constructor, e.g.: "static_friction",
+                "dynamic_friction", and "restitution"
             load_config (None or dict): If specified, should contain keyword-mapped values that are relevant for
                 loading this prim at runtime.
             abilities (None or dict): If specified, manually adds specific object states to this object. It should be
@@ -88,6 +95,10 @@ def __init__(
                 If "physical", no assistive grasping will be applied (relies on contact friction + finger force).
                 If "assisted", will magnetize any object touching and within the gripper's fingers.
                 If "sticky", will magnetize any object touching the gripper's fingers.
+            finger_static_friction (None or float): If specified, specific static friction to use for robot's fingers
+            finger_dynamic_friction (None or float): If specified, specific dynamic friction to use for robot's fingers.
+                Note: If specified, this will override any ways that are found within @link_physics_materials for any
+                robot finger gripper links
             kwargs (dict): Additional keyword arguments that are used for other super() calls from subclasses, allowing
                 for flexible compositions of various object subclasses (e.g.: Robot is USDObject + ControllableObject).
         """
@@ -128,6 +139,7 @@ def __init__(
             fixed_base=fixed_base,
             visual_only=visual_only,
             self_collisions=self_collisions,
+            link_physics_materials=link_physics_materials,
             load_config=load_config,
             abilities=abilities,
             control_freq=control_freq,
@@ -141,6 +153,8 @@ def __init__(
             proprio_obs=proprio_obs,
             sensor_config=sensor_config,
             grasping_mode=grasping_mode,
+            finger_static_friction=finger_static_friction,
+            finger_dynamic_friction=finger_dynamic_friction,
             grasping_direction=(
                 "lower" if end_effector == "gripper" else "upper"
             ),  # gripper grasps in the opposite direction

omnigibson/robots/behavior_robot.py

@@ -68,6 +68,7 @@ def __init__(
         visible=True,
         visual_only=False,
         self_collisions=True,
+        link_physics_materials=None,
         load_config=None,
         # Unique to USDObject hierarchy
         abilities=None,
@@ -84,6 +85,8 @@ def __init__(
         proprio_obs="default",
         # Unique to ManipulationRobot
         grasping_mode="assisted",
+        finger_static_friction=None,
+        finger_dynamic_friction=None,
         # unique to BehaviorRobot
         use_ghost_hands=True,
         **kwargs,
@@ -102,6 +105,7 @@ def __init__(
             fixed_base=True,
             visual_only=visual_only,
             self_collisions=self_collisions,
+            link_physics_materials=link_physics_materials,
             load_config=load_config,
             abilities=abilities,
             control_freq=control_freq,
@@ -114,6 +118,8 @@ def __init__(
             exclude_sensor_names=exclude_sensor_names,
             proprio_obs=proprio_obs,
             grasping_mode=grasping_mode,
+            finger_static_friction=finger_static_friction,
+            finger_dynamic_friction=finger_dynamic_friction,
             grasping_direction="upper",
             **kwargs,
         )