From 8b6340855b581faf9611ccad776b80dd0ef742bd Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 5 Aug 2026 19:03:54 +0800 Subject: [PATCH 1/2] fix(randomization): stabilize visual material lifecycle Preserve deterministic material assignments as reset baselines, keep non-random setters when visual randomization is filtered, and reuse bounded texture pools in the legacy path. --- embodichain/lab/gym/envs/embodied_env.py | 12 ++- .../gym/envs/managers/randomization/visual.py | 32 +++++--- embodichain/lab/sim/objects/articulation.py | 21 ++++++ embodichain/lab/sim/objects/rigid_object.py | 13 ++++ embodichain/lab/sim/sim_manager.py | 2 +- .../test_randomize_visual_material.py | 75 ++++++++++++++++++- tests/gym/envs/test_embodied_env.py | 23 +++++- .../test_asset_material_initialization.py | 21 ++++++ 8 files changed, 180 insertions(+), 19 deletions(-) diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 3c9173c45..a8e717008 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/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): diff --git a/embodichain/lab/gym/envs/managers/randomization/visual.py b/embodichain/lab/gym/envs/managers/randomization/visual.py index fbb35f210..457c3eec5 100644 --- a/embodichain/lab/gym/envs/managers/randomization/visual.py +++ b/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, diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 7b2297230..a2ef24f2b 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/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( diff --git a/embodichain/lab/sim/objects/rigid_object.py b/embodichain/lab/sim/objects/rigid_object.py index 66ae8dc26..5a4bd80a1 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/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( diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 2ee8ae064..c6e8ce9e5 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/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() diff --git a/tests/gym/envs/managers/test_randomize_visual_material.py b/tests/gym/envs/managers/test_randomize_visual_material.py index bb4edaf47..85cb3c821 100644 --- a/tests/gym/envs/managers/test_randomize_visual_material.py +++ b/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"), [ diff --git a/tests/gym/envs/test_embodied_env.py b/tests/gym/envs/test_embodied_env.py index 664456a6f..77b8334bf 100644 --- a/tests/gym/envs/test_embodied_env.py +++ b/tests/gym/envs/test_embodied_env.py @@ -17,14 +17,21 @@ from __future__ import annotations import os +from types import SimpleNamespace + import torch import pytest import numpy as np import gymnasium as gym from embodichain.lab.sim.cfg import RenderCfg -from embodichain.lab.gym.envs import EmbodiedEnvCfg +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg from embodichain.lab.sim.objects import RigidObject, Robot +from embodichain.lab.gym.envs.managers.cfg import EventCfg +from embodichain.lab.gym.envs.managers.randomization.visual import ( + randomize_visual_material, + set_rigid_object_visual_material, +) from embodichain.lab.gym.utils.gym_utils import config_to_cfg, DEFAULT_MANAGER_MODULES from embodichain.lab.gym.utils.registration import register_env from embodichain.lab.sim import SimulationManager, SimulationManagerCfg @@ -119,6 +126,20 @@ } +def test_visual_randomization_filter_keeps_deterministic_material_events(): + events = SimpleNamespace( + random_material=EventCfg(func=randomize_visual_material), + set_material=EventCfg(func=set_rigid_object_visual_material), + ) + env = EmbodiedEnv.__new__(EmbodiedEnv) + env.cfg = SimpleNamespace(filter_visual_rand=True, events=events) + + env._apply_functor_filter() + + assert events.random_material is None + assert events.set_material is not None + + class EmbodiedEnvTest: """Shared test logic for CPU and CUDA.""" diff --git a/tests/sim/objects/test_asset_material_initialization.py b/tests/sim/objects/test_asset_material_initialization.py index 5e3c242c8..6602e811b 100644 --- a/tests/sim/objects/test_asset_material_initialization.py +++ b/tests/sim/objects/test_asset_material_initialization.py @@ -53,6 +53,9 @@ def _make_asset(asset_type, materials): entity = MagicMock(name="entity") entity.get_render_body.return_value = render_body + entity.set_material.side_effect = lambda *args: materials.__setitem__( + slice(None), [args[-1]] * len(materials) + ) asset = asset_type.__new__(asset_type) asset._entities = [entity] @@ -126,6 +129,24 @@ def test_asset_restores_original_material_after_replacement(asset_type): assert asset.is_shared_visual_material is False +@pytest.mark.parametrize("asset_type", (RigidObject, Articulation)) +def test_asset_can_update_default_visual_material(asset_type): + current = [_material("original")] + asset, render_body = _make_asset(asset_type, current) + asset._initialize_existing_visual_material() + replacement = _replacement_material() + + asset.set_visual_material(replacement, update_default=True) + replacement_inst = replacement.create_instance.return_value + current[0] = _material("temporary") + render_body.set_material.reset_mock() + + asset.restore_visual_material() + + render_body.set_material.assert_called_once_with(0, replacement_inst.mat) + assert _registered_material(asset) is replacement_inst + + def test_asset_restores_empty_original_assignment(asset_type): current = [None] asset, render_body = _make_asset(asset_type, current) From 70e2b88aa1e108db16e4e64bfd0480dcd7ecc454 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 5 Aug 2026 20:03:22 +0800 Subject: [PATCH 2/2] wip --- embodichain/lab/scripts/run_env.py | 31 +++++++++++++++++++++++++++--- tests/lab/scripts/test_run_env.py | 23 ++++++++++++++++++++++ 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index 4fc5e6600..ce9444276 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/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: diff --git a/tests/lab/scripts/test_run_env.py b/tests/lab/scripts/test_run_env.py index 7950d5a56..df08e02a3 100644 --- a/tests/lab/scripts/test_run_env.py +++ b/tests/lab/scripts/test_run_env.py @@ -16,11 +16,34 @@ from __future__ import annotations +from unittest.mock import MagicMock + from embodichain.lab.gym.utils.gym_utils import merge_args_with_gym_config +from embodichain.lab.scripts import run_env from embodichain.lab.scripts.run_env import _create_parser GYM_CONFIG_PATH = "task.yaml" GYM_ID = "Dummy-v0" +EPISODE_INDEX = 3 +ACTION_LIST_INDEX = 0 + + +def test_generate_function_displays_episode_and_action_list_indices( + monkeypatch, +) -> None: + """Progress output distinguishes episodes from their local action lists.""" + env = MagicMock() + env.reset.return_value = (None, {}) + env.get_wrapper_attr.return_value.return_value = [object()] + env.step.return_value = (None, None, None, None, None) + progress = MagicMock(side_effect=lambda actions, **kwargs: actions) + monkeypatch.setattr(run_env.tqdm, "tqdm", progress) + + run_env.generate_function(env, num_traj=1, time_id=EPISODE_INDEX) + + assert progress.call_args.kwargs["desc"] == ( + f"Executing episode #{EPISODE_INDEX}, action list #{ACTION_LIST_INDEX}" + ) def test_run_env_syncs_viser_images_each_step_by_default() -> None: