Fix visual material randomization lifecycle

embodichain/lab/gym/envs/embodied_env.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 # ----------------------------------------------------------------------------
 
+from __future__ import annotations
+
 from math import log
 from functools import wraps
 from datetime import datetime
@@ -413,9 +415,13 @@ def _apply_functor_filter(self) -> None:
         from embodichain.utils.module_utils import get_all_exported_items_from_module
         from embodichain.lab.gym.envs.managers.cfg import EventCfg
 
-        functors_to_remove = get_all_exported_items_from_module(
-            "embodichain.lab.gym.envs.managers.randomization.visual"
-        )
+        functors_to_remove = {
+            name
+            for name in get_all_exported_items_from_module(
+                "embodichain.lab.gym.envs.managers.randomization.visual"
+            )
+            if name.startswith("randomize_")
+        }
         if self.cfg.filter_visual_rand and self.cfg.events:
             # Iterate through all attributes of the events object
             for attr_name in dir(self.cfg.events):

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

@@ -100,7 +100,7 @@ def set_rigid_object_visual_material(
 
     mat = env.sim.create_visual_material(mat_cfg)
     obj: RigidObject = env.sim.get_rigid_object(entity_cfg.uid)
-    obj.set_visual_material(mat, env_ids=env_ids)
+    obj.set_visual_material(mat, env_ids=env_ids, update_default=True)
 
 
 def set_rigid_object_group_visual_material(
@@ -645,6 +645,14 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv):
                 self._new_mode = False
         if not self._new_mode:
             self._init_legacy(env)
+            # The legacy path used to upload a new Texture on every invocation.
+            # This is especially easy to hit for the default plane and for an
+            # automatic reuse-to-legacy fallback, where clean_materials() is not
+            # safe to call because live Python Material objects are retained.
+            # Reuse the same bounded texture pools as the existing-material path
+            # instead, so repeated randomization does not consume texture IDs.
+            self._build_library_textures(env)
+            self._build_solid_textures(env)
 
     def _init_reuse(self, env: EmbodiedEnv) -> None:
         """Init the reuse path: capture existing materials, pre-create textures, resolve tiers."""
@@ -810,10 +818,13 @@ def gen_random_base_color_texture(width: int, height: int) -> torch.Tensor:
         return rgba
 
     def _randomize_texture(self, mat_inst: VisualMaterialInst) -> None:
-        if len(self.textures) > 0:
-            # Randomly select a texture from the preloaded textures
-            texture_idx = torch.randint(0, len(self.textures), (1,)).item()
-            mat_inst.set_base_color_texture(texture_data=self.textures[texture_idx])
+        if self._library_textures:
+            # Bind a pre-created Texture instead of uploading the same image on
+            # every randomization interval.
+            texture_idx = torch.randint(0, len(self._library_textures), (1,)).item()
+            mat_inst.set_base_color_texture(
+                texture_obj=self._library_textures[texture_idx]
+            )
 
     def _randomize_mat_inst(
         self,
@@ -823,7 +834,7 @@ def _randomize_mat_inst(
         idx: int = 0,
     ) -> None:
         # randomize texture or base color based on the probability.
-        if random.random() < random_texture_prob and len(self.textures) != 0:
+        if random.random() < random_texture_prob and self._library_textures:
             for key, value in plan.items():
                 if key == "base_color":
                     mat_inst.set_base_color(value[idx].tolist())
@@ -832,11 +843,12 @@ def _randomize_mat_inst(
 
             self._randomize_texture(mat_inst)
         else:
-            # set a random base color instead.
-            random_color_texture = (
-                randomize_visual_material.gen_random_base_color_texture(2, 2)
+            # Use the bounded solid-color palette. Uploading a generated tensor
+            # here would allocate a fresh DexSim texture ID on every call.
+            texture_idx = torch.randint(0, len(self._solid_textures), (1,)).item()
+            mat_inst.set_base_color_texture(
+                texture_obj=self._solid_textures[texture_idx]
             )
-            mat_inst.set_base_color_texture(texture_data=random_color_texture)
 
     def __call__(
         self,

embodichain/lab/scripts/run_env.py

@@ -42,7 +42,26 @@
 from embodichain.utils.logger import log_warning, log_info, log_error
 
 
-def generate_and_execute_action_list(env, idx, debug_mode, **kwargs):
+def generate_and_execute_action_list(
+    env: gymnasium.Env,
+    idx: int,
+    debug_mode: bool,
+    *,
+    episode_idx: int = 0,
+    **kwargs: object,
+) -> bool:
+    """Generate and execute one demonstration action list.
+
+    Args:
+        env: Environment used to generate and execute the actions.
+        idx: Index of the action list within the current episode.
+        debug_mode: Whether debug mode is enabled.
+        episode_idx: Index of the current episode.
+        **kwargs: Additional arguments forwarded to action generation.
+
+    Returns:
+        Whether a non-empty action list was generated and executed.
+    """
 
     action_list = env.get_wrapper_attr("create_demo_action_list")(
         action_sentence=idx, **kwargs
@@ -53,7 +72,9 @@ def generate_and_execute_action_list(env, idx, debug_mode, **kwargs):
         return False
 
     for action in tqdm.tqdm(
-        action_list, desc=f"Executing action list #{idx}", unit="step"
+        action_list,
+        desc=f"Executing episode #{episode_idx}, action list #{idx}",
+        unit="step",
     ):
         # Step the environment with the current action
         # The environment will automatically detect truncation based on action_length
@@ -99,7 +120,11 @@ def generate_function(
         ret = []
         for trajectory_idx in range(num_traj):
             valid = generate_and_execute_action_list(
-                env, trajectory_idx, debug_mode, **kwargs
+                env,
+                trajectory_idx,
+                debug_mode,
+                episode_idx=time_id,
+                **kwargs,
             )
 
             if not valid:

embodichain/lab/sim/objects/articulation.py

@@ -2175,6 +2175,7 @@ def set_visual_material(
         env_ids: Sequence[int] | None = None,
         link_names: List[str] | None = None,
         shared: bool = False,
+        update_default: bool = False,
     ) -> None:
         """Set visual material for the rigid object.
 
@@ -2183,6 +2184,8 @@ def set_visual_material(
             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.
+            update_default: Whether the assigned material should become the baseline
+                restored by :meth:`reset`. 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
@@ -2196,6 +2199,15 @@ def set_visual_material(
                 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
+                    if update_default:
+                        self._original_visual_material[env_idx][link_name] = (
+                            _capture_render_materials(
+                                self._entities[env_idx].get_render_body(link_name)
+                            )
+                        )
+                        self._original_visual_material_inst[env_idx][
+                            link_name
+                        ] = mat_inst
             self.is_shared_visual_material = True
         else:
             for i, env_idx in enumerate(local_env_ids):
@@ -2205,6 +2217,15 @@ def set_visual_material(
                     )
                     self._entities[env_idx].set_material(link_name, mat_inst.mat)
                     self._visual_material[env_idx][link_name] = mat_inst
+                    if update_default:
+                        self._original_visual_material[env_idx][link_name] = (
+                            _capture_render_materials(
+                                self._entities[env_idx].get_render_body(link_name)
+                            )
+                        )
+                        self._original_visual_material_inst[env_idx][
+                            link_name
+                        ] = mat_inst
             self.is_shared_visual_material = False
 
     def get_visual_material_inst(

embodichain/lab/sim/objects/rigid_object.py

@@ -863,6 +863,7 @@ def set_visual_material(
         mat: VisualMaterial,
         env_ids: Sequence[int] | None = None,
         shared: bool = False,
+        update_default: bool = False,
     ) -> None:
         """Set visual material for the rigid object.
 
@@ -874,6 +875,8 @@ def set_visual_material(
             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.
+            update_default: Whether the assigned material should become the baseline
+                restored by :meth:`reset`. Defaults to False.
         """
         local_env_ids = self._all_indices if env_ids is None else env_ids
 
@@ -885,12 +888,22 @@ def set_visual_material(
             for env_idx in local_env_ids:
                 self._entities[env_idx].set_material(mat_inst.mat)
                 self._visual_material[env_idx] = mat_inst
+                if update_default:
+                    self._original_visual_material[env_idx] = _capture_render_materials(
+                        self._entities[env_idx].get_render_body()
+                    )
+                    self._original_visual_material_inst[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
+                if update_default:
+                    self._original_visual_material[env_idx] = _capture_render_materials(
+                        self._entities[env_idx].get_render_body()
+                    )
+                    self._original_visual_material_inst[env_idx] = mat_inst
             self.is_shared_visual_material = False
 
     def get_visual_material_inst(

embodichain/lab/sim/sim_manager.py

@@ -1242,7 +1242,7 @@ def add_rigid_object(
 
         if cfg.shape.visual_material:
             mat = self.create_visual_material(cfg.shape.visual_material)
-            rigid_obj.set_visual_material(mat)
+            rigid_obj.set_visual_material(mat, update_default=True)
 
         self._rigid_objects[uid] = rigid_obj
         self.notify_visualization_topology_changed()

tests/gym/envs/managers/test_randomize_visual_material.py

@@ -27,8 +27,13 @@
 from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg
 from embodichain.lab.gym.envs.managers.randomization.visual import (
     randomize_visual_material,
+    set_rigid_object_visual_material,
+)
+from embodichain.lab.sim.material import (
+    ReuseSegmentState,
+    VisualMaterialCfg,
+    VisualMaterialInst,
 )
-from embodichain.lab.sim.material import ReuseSegmentState, VisualMaterialInst
 from embodichain.lab.sim.objects.articulation import Articulation
 from embodichain.lab.sim.objects.rigid_object import RigidObject
 
@@ -126,6 +131,25 @@ def _segment(mesh_id: int = 0, original=None) -> ReuseSegmentState:
     )
 
 
+def test_deterministic_material_setter_updates_reset_baseline():
+    env = _MockEnv()
+    obj = env.sim.get_asset("obj")
+    material = MagicMock(name="material")
+    env.sim.create_visual_material = MagicMock(return_value=material)
+    env.sim.get_rigid_object_uid_list = MagicMock(return_value=["obj"])
+    env.sim.get_rigid_object = MagicMock(return_value=obj)
+
+    set_rigid_object_visual_material(
+        env,
+        None,
+        SceneEntityCfg(uid="obj"),
+        VisualMaterialCfg(uid="fixed"),
+    )
+
+    _, kwargs = obj.set_visual_material.call_args
+    assert kwargs["update_default"] is True
+
+
 def _make_rigid_functor(
     params: dict | None = None,
     *,
@@ -196,6 +220,26 @@ def test_reuse_init_degrades_to_legacy_on_failure():
     assert env.sim.created_visual_materials == ["obj_random_mat"]
 
 
+def test_automatic_legacy_fallback_reuses_bounded_texture_pool():
+    env = _MockEnv()
+    env.sim.get_asset("obj").get_existing_visual_material.side_effect = ValueError(
+        "no material"
+    )
+    palette_size = 2
+    functor = randomize_visual_material(
+        _make_cfg({"random_texture_prob": 0.0, "solid_texture_count": palette_size}),
+        env,
+    )
+    created_at_init = env.sim.env.create_color_texture.call_count
+
+    for _ in range(1025):
+        _run(functor, env)
+
+    assert created_at_init == palette_size
+    assert env.sim.env.create_color_texture.call_count == created_at_init
+    env.sim.env.clean_materials.assert_not_called()
+
+
 def test_reuse_call_reattaches_without_cleaning():
     env, obj, functor = _make_rigid_functor()
     _force_tier(functor, tier=2)
@@ -245,18 +289,41 @@ def test_articulation_samples_tier_per_link():
     assert solid_call.kwargs["texture_obj"] in functor._solid_textures
 
 
-def test_default_plane_randomizes_in_place_without_cleaning():
+def test_default_plane_reuses_bounded_texture_pool_without_cleaning():
     env = _MockEnv()
-    functor = randomize_visual_material(_make_cfg(uid="default_plane"), env)
+    palette_size = 2
+    functor = randomize_visual_material(
+        _make_cfg(
+            {"random_texture_prob": 0.0, "solid_texture_count": palette_size},
+            uid="default_plane",
+        ),
+        env,
+    )
+    created_at_init = env.sim.env.create_color_texture.call_count
     env.sim.env.clean_materials.reset_mock()
 
-    _run(functor, env)
+    for _ in range(1025):
+        _run(functor, env)
 
     assert functor._new_mode is False
     assert env.sim.created_visual_materials == []
+    assert created_at_init == palette_size
+    assert env.sim.env.create_color_texture.call_count == created_at_init
     env.sim.env.clean_materials.assert_not_called()
 
 
+def test_legacy_library_randomization_binds_precreated_texture():
+    env = _MockEnv()
+    functor = randomize_visual_material(_make_cfg({"fallback_to_new": True}), env)
+    texture = MagicMock(name="library_texture")
+    functor._library_textures = [texture]
+    mat_inst = MagicMock(spec=VisualMaterialInst)
+
+    functor._randomize_mat_inst(mat_inst, {}, random_texture_prob=1.0)
+
+    mat_inst.set_base_color_texture.assert_called_once_with(texture_obj=texture)
+
+
 @pytest.mark.parametrize(
     ("params", "has_library", "expected"),
     [