Support shared visual material instances (#110)

Author: yuecideng

embodichain/lab/gym/envs/managers/randomization/visual.py

@@ -534,6 +534,24 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv):
                 self.entity_cfg.link_names = link_names
                 self.entity.set_visual_material(mat, link_names=link_names)
 
+        # ground plane only has one instance.
+        self._mat_insts = None
+        if self.entity_cfg.uid == "default_plane":
+            self._mat_insts = env.sim.get_visual_material(
+                "plane_mat"
+            ).get_default_instance()
+            return
+        elif isinstance(self.entity, RigidObject):
+            self._mat_insts = self.entity.get_visual_material_inst()
+            if self.entity.is_shared_visual_material:
+                self._mat_insts = self._mat_insts[:1]
+        elif isinstance(self.entity, Articulation):
+            self._mat_insts = self.entity.get_visual_material_inst(
+                link_names=self.entity_cfg.link_names,
+            )
+            if self.entity.is_shared_visual_material:
+                self._mat_insts = self._mat_insts[:1]
+
     @staticmethod
     def gen_random_base_color_texture(width: int, height: int) -> torch.Tensor:
         """Generate a random base color texture.
@@ -643,8 +661,6 @@ def __call__(
             )
             randomize_plan["ior"] = ior
 
-        # ground plane only has one instance.
-        mat_insts = None
         if self.entity_cfg.uid == "default_plane":
             mat_inst = env.sim.get_visual_material("plane_mat").get_default_instance()
             self._randomize_mat_inst(
@@ -654,15 +670,8 @@ def __call__(
                 idx=0,
             )
             return
-        elif isinstance(self.entity, RigidObject):
-            mat_insts = self.entity.get_visual_material_inst(env_ids=env_ids)
-        elif isinstance(self.entity, Articulation):
-            mat_insts = self.entity.get_visual_material_inst(
-                env_ids=env_ids,
-                link_names=self.entity_cfg.link_names,
-            )
 
-        for i, data in enumerate(mat_insts):
+        for i, data in enumerate(self._mat_insts):
             if isinstance(self.entity, RigidObject):
                 # For RigidObject, data is the material instance directly
                 mat: VisualMaterialInst = data

embodichain/lab/sim/objects/articulation.py

@@ -628,6 +628,7 @@ def __init__(
         self._visual_material: List[Dict[str, VisualMaterialInst]] = [
             {} for _ in range(len(entities))
         ]
+        self.is_shared_visual_material = False
 
         # Stores mimic information for joints.
         self._mimic_info = entities[0].get_mimic_info()
@@ -1582,24 +1583,38 @@ def set_visual_material(
         mat: VisualMaterial,
         env_ids: Sequence[int] | None = None,
         link_names: List[str] | None = None,
+        shared: bool = False,
     ) -> None:
         """Set visual material for the rigid object.
 
         Args:
             mat (VisualMaterial): The material to set.
             env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used.
             link_names (List[str] | None, optional): List of link names to apply the material to. If None, applies to all links.
+            shared (bool, optional): Whether to share the material instance across links and environments. Defaults to False.
         """
         local_env_ids = self._all_indices if env_ids is None else env_ids
         link_names = self.link_names if link_names is None else link_names
 
-        for i, env_idx in enumerate(local_env_ids):
+        if shared:
+            if len(local_env_ids) != self.num_instances:
+                logger.log_error(f"Cannot share material instance for partial env_ids.")
+
             for link_name in link_names:
-                mat_inst = mat.create_instance(
-                    f"{mat.uid}_{self.uid}_{link_name}_{env_idx}"
-                )
-                self._entities[env_idx].set_material(link_name, mat_inst.mat)
-                self._visual_material[env_idx][link_name] = mat_inst
+                mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}_{link_name}")
+                for i, env_idx in enumerate(local_env_ids):
+                    self._entities[env_idx].set_material(link_name, mat_inst.mat)
+                    self._visual_material[env_idx][link_name] = mat_inst
+            self.is_shared_visual_material = True
+        else:
+            for i, env_idx in enumerate(local_env_ids):
+                for link_name in link_names:
+                    mat_inst = mat.create_instance(
+                        f"{mat.uid}_{self.uid}_{link_name}_{env_idx}"
+                    )
+                    self._entities[env_idx].set_material(link_name, mat_inst.mat)
+                    self._visual_material[env_idx][link_name] = mat_inst
+            self.is_shared_visual_material = False
 
     def get_visual_material_inst(
         self,

embodichain/lab/sim/objects/rigid_object.py

@@ -209,6 +209,7 @@ def __init__(
 
         # For rendering purposes, each instance can have its own material.
         self._visual_material: List[VisualMaterialInst] = [None] * len(entities)
+        self.is_shared_visual_material = False
 
         for entity in entities:
             entity.set_body_scale(*cfg.body_scale)
@@ -566,20 +567,39 @@ def get_mass(self, env_ids: Sequence[int] | None = None) -> torch.Tensor:
         return torch.as_tensor(masses, dtype=torch.float32, device=self.device)
 
     def set_visual_material(
-        self, mat: VisualMaterial, env_ids: Sequence[int] | None = None
+        self,
+        mat: VisualMaterial,
+        env_ids: Sequence[int] | None = None,
+        shared: bool = False,
     ) -> None:
         """Set visual material for the rigid object.
 
+        Note:
+            If `shared` is True, the same material instance will be used for all specified environment indices.
+            If `shared` is False, a unique material instance will be created for each specified environment index.
+
         Args:
             mat (VisualMaterial): The material to set.
             env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used.
+            shared (bool, optional): Whether to share the material instance among all specified environment indices. Defaults to False.
         """
         local_env_ids = self._all_indices if env_ids is None else env_ids
 
-        for i, env_idx in enumerate(local_env_ids):
-            mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}_{env_idx}")
-            self._entities[env_idx].set_material(mat_inst.mat)
-            self._visual_material[env_idx] = mat_inst
+        if shared:
+            if len(local_env_ids) != self.num_instances:
+                logger.log_error(f"Cannot share material instance for partial env_ids.")
+
+            mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}")
+            for env_idx in local_env_ids:
+                self._entities[env_idx].set_material(mat_inst.mat)
+                self._visual_material[env_idx] = mat_inst
+            self.is_shared_visual_material = True
+        else:
+            for i, env_idx in enumerate(local_env_ids):
+                mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}_{env_idx}")
+                self._entities[env_idx].set_material(mat_inst.mat)
+                self._visual_material[env_idx] = mat_inst
+            self.is_shared_visual_material = False
 
     def get_visual_material_inst(
         self, env_ids: Sequence[int] | None = None