diff --git a/docs/source/guides/custom_functors.md b/docs/source/guides/custom_functors.md index be48fd314..7f2430865 100644 --- a/docs/source/guides/custom_functors.md +++ b/docs/source/guides/custom_functors.md @@ -4,6 +4,37 @@ Functors are the building blocks of EmbodiChain's manager system. They define ho This guide explains the two functor styles (function and class), how to register them in manager configs, and provides examples for each functor type. +## Per-instance visual texture randomization + +The `randomize_visual_material` event functor can assign base-color textures independently to vectorized environment instances. Add these options to its `params` mapping: + +| Option | Values | Purpose | +|---|---|---| +| `texture_sampling` | `random`, `without_replacement`, `cycle`, `fixed` | Selects how textures are assigned to target environments. | +| `texture_indices` | `dict[int, int]` | Maps global environment IDs to texture indices when using `fixed`. | +| `texture_scope` | `per_material`, `per_instance` | For articulations, choose per-link textures or one texture shared by all selected links in each environment. | + +For unique assignments in a reset batch, use `without_replacement` (there must be at least as many source textures as target environments): + +```yaml +params: + entity_cfg: {uid: bottle} + texture_path: DexsimMaterials/Bottle + texture_sampling: without_replacement + texture_scope: per_instance +``` + +For deterministic assignments, use `fixed` and provide an index for every environment that may be reset: + +```yaml +params: + entity_cfg: {uid: bottle} + texture_path: DexsimMaterials/Bottle + texture_sampling: fixed + texture_indices: {0: 2, 1: 0, 2: 3, 3: 1} + texture_scope: per_instance +``` + --- ## Functor Basics diff --git a/docs/superpowers/plans/2026-07-14-per-instance-visual-texture-randomization-plan.md b/docs/superpowers/plans/2026-07-14-per-instance-visual-texture-randomization-plan.md new file mode 100644 index 000000000..49c69344b --- /dev/null +++ b/docs/superpowers/plans/2026-07-14-per-instance-visual-texture-randomization-plan.md @@ -0,0 +1,184 @@ +# Per-Instance Visual Texture Randomization Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (\`- [ ]\`) syntax for tracking. + +**Goal:** Make \`randomize_visual_material\` safely assign random, unique, or fixed base-color textures to selected object instances across vectorized environments. + +**Architecture:** Keep the class-style functor and per-environment material instances. Add a cached DexSim texture-reference path to \`VisualMaterialInst\`, select textures by global environment ID, retrieve only targeted instances, and remove unsafe runtime cleanup. + +**Tech Stack:** Python 3.11+, PyTorch, DexSim 0.4.3, pytest, existing \`Functor\`/\`VisualMaterialInst\` APIs. + +## Global Constraints + +- Preserve default random-with-replacement behavior. +- Use full annotations and \`from __future__ import annotations\` for new Python APIs. +- Preserve Apache headers on modified Python files. +- Never call \`env.clean_materials()\` during an episode. +- Do not modify the unrelated working-tree change in \`embodichain/gen_sim/action_agent_pipeline/generation/action_agent_config_defaults.yaml\`. +- \`RigidObjectGroup\` remains out of scope. + +--- + +## File Map + +- Modify \`embodichain/lab/sim/material.py\`: accept a pre-created DexSim texture reference. +- Modify \`embodichain/lab/gym/envs/managers/randomization/visual.py\`: normalize IDs, cache references, select textures, and apply only to targets. +- Create \`tests/gym/envs/managers/test_randomize_visual_material.py\`: fake-backed unit tests and simulator integration coverage. + +## Task 1: Add the texture-reference material API + +**Files:** \`embodichain/lab/sim/material.py:244-271\`; new manager test file. + +**Produces:** \`VisualMaterialInst.set_base_color_texture(texture_ref=...)\`, storing the reference and calling \`set_base_color_map(texture_ref)\` without creating another texture. + +- [ ] **Step 1: Write the failing test** + +\`\`\`python +def test_set_base_color_texture_uses_texture_ref(): + material = FakeMaterial() + instance = VisualMaterialInst("instance", material) + texture_ref = object() + instance.set_base_color_texture(texture_ref=texture_ref) + assert material.get_inst("instance").base_color_map is texture_ref + assert instance.base_color_texture is texture_ref +\`\`\` + +- [ ] **Step 2: Run and verify failure** + +\`\`\`bash +pytest tests/gym/envs/managers/test_randomize_visual_material.py::test_set_base_color_texture_uses_texture_ref -q +\`\`\` + +Expected: failure because \`texture_ref\` is not accepted. + +- [ ] **Step 3: Implement the minimal branch** + +Add \`texture_ref: object | None = None\`; reject or warn if combined with \`texture_path\` or \`texture_data\`; assign \`self.base_color_texture = texture_ref\` and call the DexSim material-instance setter. + +- [ ] **Step 4: Verify** + +\`\`\`bash +pytest tests/gym/envs/managers/test_randomize_visual_material.py::test_set_base_color_texture_uses_texture_ref -q +pytest tests/sim/objects/test_rigid_object.py -q +\`\`\` + +Expected: both commands pass. + +## Task 2: Add target normalization and texture selection + +**Files:** \`embodichain/lab/gym/envs/managers/randomization/visual.py\`; new manager test file. + +**Produces:** \`_normalize_env_ids(env_ids, num_envs) -> list[int]\` and \`_select_texture_indices(mode, target_ids, num_textures, texture_indices) -> list[int]\`. + +- [ ] **Step 1: Write failing tests** + +\`\`\`python +def test_normalize_env_ids_supports_all_input_forms(): + assert _normalize_env_ids(None, 3) == [0, 1, 2] + assert _normalize_env_ids(torch.tensor([2, 0]), 3) == [2, 0] + assert _normalize_env_ids([1], 3) == [1] + assert _normalize_env_ids(slice(None), 3) == [0, 1, 2] + +def test_texture_selection_modes(): + assert sorted(_select_texture_indices("without_replacement", [0, 1, 2], 3, None)) == [0, 1, 2] + assert _select_texture_indices("fixed", [3, 1], 4, {1: 2, 3: 0}) == [0, 2] + with pytest.raises(ValueError, match="without_replacement"): + _select_texture_indices("without_replacement", [0, 1], 1, None) +\`\`\` + +- [ ] **Step 2: Run and verify failure** + +\`\`\`bash +pytest tests/gym/envs/managers/test_randomize_visual_material.py -k "normalize or selection" -q +\`\`\` + +Expected: failure because the helpers do not exist. + +- [ ] **Step 3: Implement helpers** + +Handle \`None\`, tensors, Python sequences, and \`slice(None)\`; validate IDs. Implement \`random\` with replacement, \`without_replacement\` with \`torch.randperm\`, \`cycle\` modulo source count, and \`fixed\` using global \`env_id -> texture_index\`. Raise \`ValueError\` for unknown modes, missing mappings, missing IDs, insufficient sources, or invalid indices. + +- [ ] **Step 4: Verify and commit** + +\`\`\`bash +pytest tests/gym/envs/managers/test_randomize_visual_material.py -k "normalize or selection" -q +git add embodichain/lab/gym/envs/managers/randomization/visual.py tests/gym/envs/managers/test_randomize_visual_material.py +git commit -m "test: define per-instance texture selection behavior" +\`\`\` + +Expected: focused tests pass and the commit contains only Task 2 files. + +## Task 3: Integrate targeted assignment into the functor + +**Files:** \`embodichain/lab/gym/envs/managers/randomization/visual.py:538-798\`; new manager test file. + +**Interface:** Extend \`__call__\` with \`texture_sampling: str = "random"\`, \`texture_indices: Mapping[int, int] | None = None\`, and \`texture_scope: str = "per_material"\`. + +- [ ] **Step 1: Write failing tests** + +Use fake rigid objects and material instances. Call with \`env_ids=torch.tensor([1, 3])\`; assert only material instances 1 and 3 change. Add a fixed mapping test for \`{1: 2, 3: 0}\` and an articulation \`per_instance\` test asserting all selected links in one environment share one selected texture. + +- [ ] **Step 2: Run and verify failure** + +\`\`\`bash +pytest tests/gym/envs/managers/test_randomize_visual_material.py -k "partial or fixed_assignment or per_instance" -q +\`\`\` + +Expected: failure because the current functor loops its full cached list and has no selection parameters. + +- [ ] **Step 3: Implement** + +Normalize IDs at call time; build all plans with \`len(target_ids)\` rows; select texture indices by global IDs; retrieve \`RigidObject.get_visual_material_inst(env_ids=target_ids)\` or the articulation equivalent; apply each row to its target. For \`per_instance\`, reuse one index across selected links. Apply metallic, roughness, and IOR in both texture and generated-color branches. Remove the final \`env.clean_materials()\` call. + +- [ ] **Step 4: Verify** + +\`\`\`bash +pytest tests/gym/envs/managers/test_randomize_visual_material.py -q +pytest tests/gym/envs/managers -q +\`\`\` + +Expected: all focused and manager tests pass. + +## Task 4: Cache texture references and add integration coverage + +**Files:** \`embodichain/lab/gym/envs/managers/randomization/visual.py:589-607\`; new manager test file. + +- [ ] **Step 1: Write failing cache test** + +Fake \`create_color_texture\` with a counter; invoke the functor twice and assert each source image is converted exactly once. + +- [ ] **Step 2: Run and verify failure** + +\`\`\`bash +pytest tests/gym/envs/managers/test_randomize_visual_material.py::test_texture_references_are_created_once -q +\`\`\` + +Expected: failure because current assignment recreates a GPU texture from image data each time. + +- [ ] **Step 3: Implement and test** + +Convert each padded RGBA image once during initialization, store \`self.texture_refs\`, and pass the selected reference through \`_randomize_texture\` and \`_randomize_mat_inst\`. Use canonical resolved source paths as cache keys. Add real four-environment tests for distinct no-replacement assignments, fixed mapping, partial-reset isolation, and a second reset with valid handles. + +\`\`\`bash +pytest tests/gym/envs/managers/test_randomize_visual_material.py -q +pytest tests/sim/objects/test_rigid_object.py -q +pytest tests/sim/objects/test_articulation.py -q +\`\`\` + +Expected: all selected tests pass; if renderer setup is unavailable, report that limitation while retaining pure tests. + +- [ ] **Step 4: Commit** + +\`\`\`bash +git add embodichain/lab/gym/envs/managers/randomization/visual.py embodichain/lab/sim/material.py tests/gym/envs/managers/test_randomize_visual_material.py +git commit -m "feat: support per-instance visual texture randomization" +\`\`\` + +## Task 5: Final verification and documentation + +**Files:** the relevant existing randomization documentation page, if present; no unrelated files. + +- [ ] **Step 1:** Document \`texture_sampling\`, \`texture_indices\`, and \`texture_scope\` with the fixed and no-replacement examples from the specification. +- [ ] **Step 2:** Run \`black --check --diff --color ./\` and \`git diff --check HEAD^ HEAD\`; both must pass. +- [ ] **Step 3:** Run \`pytest tests/gym/envs/managers -q\` and \`pytest tests/gym tests/sim/objects -q\`; no new failures may be attributable to this feature. +- [ ] **Step 4:** Review \`git status --short\`, \`git diff HEAD^ --stat\`, and the final feature diff; preserve the pre-existing unrelated user change. diff --git a/docs/superpowers/specs/2026-07-14-per-instance-visual-texture-randomization-design.md b/docs/superpowers/specs/2026-07-14-per-instance-visual-texture-randomization-design.md new file mode 100644 index 000000000..2e5bd3cba --- /dev/null +++ b/docs/superpowers/specs/2026-07-14-per-instance-visual-texture-randomization-design.md @@ -0,0 +1,191 @@ +# Per-Instance Visual Texture Randomization Design + +**Date:** 2026-07-14 + +**Status:** Approved design; awaiting written-spec review + +## 1. Overview + +Extend `randomize_visual_material` so each replica of a rigid object or +articulation can receive a reliably different base-color texture when +`num_envs > 1`. The extension must also support deterministic assignment by +environment ID and correctly target only the environments being reset. + +The current object wrappers already create one `VisualMaterialInst` per +environment. The existing functor randomly selects a texture for each material +instance during full-batch calls, but it does not guarantee uniqueness or offer +an explicit per-environment mapping. Its cached full material list is also +incompatible with partial resets, and it calls DexSim's unsafe runtime material +cleanup API after every application. + +## 2. Goals and Non-Goals + +### Goals + +1. Select a base-color texture independently for each targeted environment. +2. Guarantee distinct source textures for a reset when requested and enough + textures are available. +3. Permit deterministic texture assignment by global environment ID. +4. Preserve existing random sampling with replacement as the default behavior. +5. Correctly support full resets, partial resets, and interval events. +6. Reuse material and texture resources across resets without invalidating live + DexSim handles. + +### Non-Goals + +1. Add randomization of metallic, roughness, normal, or AO texture maps. +2. Support `RigidObjectGroup`; that requires a separately stored and retrievable + per-environment material-instance collection. +3. Change the default plane's global-material behavior. + +## 3. Public Configuration Interface + +Keep `randomize_visual_material` as the public functor and retain all existing +parameters. Add the following optional parameters: + +| Parameter | Type | Default | Meaning | +|---|---|---|---| +| `texture_sampling` | `"random" \| "without_replacement" \| "cycle" \| "fixed"` | `"random"` | How folder textures are chosen. | +| `texture_indices` | `dict[int, int] \| None` | `None` | Global `env_id -> texture_index` mapping; required for `"fixed"`. | +| `texture_scope` | `"per_material" \| "per_instance"` | `"per_material"` | Whether an articulation samples one texture per selected link or one texture shared by all selected links in an environment. | + +`random` retains current semantics: each target samples independently with +replacement, so duplicate textures are valid. `without_replacement` samples a +permutation and raises `ValueError` when fewer textures than targeted +environments exist. `cycle` assigns textures in source order modulo the source +count. `fixed` looks up every target environment in `texture_indices` and +raises `ValueError` for a missing or out-of-range entry. + +For the requested behavior, configure: + +```yaml +random_material: + func: randomize_visual_material + mode: reset + params: + entity_cfg: {uid: bottle} + texture_path: DexsimMaterials/Bottle + random_texture_prob: 1.0 + texture_sampling: without_replacement + texture_scope: per_instance +``` + +For a deterministic mapping: + +```yaml +random_material: + func: randomize_visual_material + mode: reset + params: + entity_cfg: {uid: bottle} + texture_path: DexsimMaterials/Bottle + texture_sampling: fixed + texture_indices: {0: 2, 1: 0, 2: 3, 3: 1} + texture_scope: per_instance +``` + +## 4. Architecture and Runtime Flow + +### 4.1 Persistent material and texture ownership + +At initialization, the functor creates one PBR material template and assigns a +unique material instance to every rigid-object replica or every selected +articulation link. The functor retains the template and the loaded texture +sources for its lifetime. + +Textures are decoded once and converted to DexSim texture references once. The +cache key uses the resolved directory path rather than just its basename to +avoid collisions between unrelated texture folders with the same name. + +`VisualMaterialInst.set_base_color_texture` gains a texture-reference input so +assignment can directly call DexSim's material-instance map setter without +creating a new GPU texture on every reset. + +### 4.2 Targeted material retrieval + +On each call, normalize `env_ids` to a CPU list of global environment IDs. +Accept `None`, tensors, Python sequences, and `slice(None)`. Return immediately +for an empty target set. + +Retrieve material instances for exactly those IDs at call time: + +- `RigidObject.get_visual_material_inst(env_ids=target_ids)` +- `Articulation.get_visual_material_inst(env_ids=target_ids, link_names=...)` + +Do not iterate a material list cached for all environments. This ensures a +partial reset affects only the selected replicas and aligns every sampled plan +row with its material instance. + +### 4.3 Sampling and application + +1. Build a material-property plan of shape `(num_targets, ...)` for base color, + metallic, roughness, and IOR. +2. Build texture indices keyed by the selected global environment IDs according + to `texture_sampling`. +3. Apply scalar PBR properties for every target, independent of whether the + texture branch or generated-color branch is taken. +4. If `random_texture_prob` selects the folder-texture branch, assign the + cached texture reference indicated by the plan. +5. Otherwise, assign a generated solid-color texture. Generated colors remain + per material for backward compatibility; `per_instance` reuses the same + generated texture across selected articulation links in an environment. + +The functor must never call `env.clean_materials()` or +`SimulationManager.clean_materials()` during an episode. DexSim documents that +the lower-level call invalidates live Python-side material and mesh references. +Simulation teardown remains responsible for cleanup. + +## 5. Error Handling + +- `texture_sampling` is unknown: raise `ValueError`. +- `fixed` without `texture_indices`: raise `ValueError`. +- A fixed mapping lacks a targeted global environment ID: raise `ValueError`. +- A texture index is negative or outside the loaded texture collection: raise + `ValueError`. +- `without_replacement` has fewer source textures than target environments: + raise `ValueError`. +- A texture-based strategy has no readable source textures: raise `ValueError`. +- If `random_texture_prob` is outside `[0, 1]`, raise `ValueError`. + +## 6. Testing Plan + +### Unit tests + +1. Normalize `None`, tensor, sequence, and `slice(None)` environment IDs. +2. Verify `random`, `without_replacement`, `cycle`, and `fixed` index plans. +3. Verify validation errors for invalid modes, mappings, source counts, and + index ranges. +4. Verify a partial target, for example `[1, 3]` of four environments, updates + only instances 1 and 3. +5. Verify property values are applied whether a folder texture or a generated + color texture is selected. + +### Simulator integration tests + +1. Create a four-environment rigid object with four distinct textures and use + `without_replacement`; assert the assigned texture identities are distinct. +2. Apply `fixed` assignment to a partial reset and assert non-target material + maps are unchanged. +3. Perform a second reset after a full reset; assert materials remain valid and + can be updated. +4. For an articulation with multiple selected links, verify `per_instance` + applies the same texture to each selected link within one environment and a + different texture in another environment. +5. Confirm cached texture creation happens only during functor initialization, + not on every reset. + +## 7. Files Expected to Change + +- `embodichain/lab/gym/envs/managers/randomization/visual.py` +- `embodichain/lab/sim/material.py` +- New focused tests under `tests/gym/envs/managers/` +- Configuration/API documentation for the functor, if a suitable randomization + reference page exists + +## 8. References + +- `embodichain/lab/gym/envs/managers/randomization/visual.py` +- `embodichain/lab/sim/material.py` +- `embodichain/lab/sim/objects/rigid_object.py` +- `embodichain/lab/sim/objects/articulation.py` +- `embodichain/lab/gym/envs/managers/event_manager.py` diff --git a/embodichain/lab/gym/envs/managers/randomization/visual.py b/embodichain/lab/gym/envs/managers/randomization/visual.py index bb4a08ac9..cd2bca878 100644 --- a/embodichain/lab/gym/envs/managers/randomization/visual.py +++ b/embodichain/lab/gym/envs/managers/randomization/visual.py @@ -23,7 +23,7 @@ import numpy as np from pathlib import Path -from typing import TYPE_CHECKING, Literal, Union, Dict +from typing import TYPE_CHECKING, Literal, Union, Dict, Sequence, Mapping from embodichain.lab.sim.objects import ( Light, @@ -64,6 +64,93 @@ ] +def _get_texture_refs(sim, key: str, textures: list[torch.Tensor]) -> list[object]: + """Return simulation-cached GPU texture references for source images.""" + cached = sim.get_texture_ref_cache(key) + if cached is not None: + return cached + import dexsim + + dex_env = dexsim.default_world().get_env() + refs = [ + dex_env.create_color_texture(image.cpu().numpy(), has_alpha=True) + for image in textures + ] + sim.set_texture_ref_cache(key, refs) + return refs + + +def _normalize_env_ids( + env_ids: torch.Tensor | Sequence[int] | slice | None, num_envs: int +) -> list[int]: + """Normalize environment selectors to validated integer IDs.""" + if num_envs < 0: + raise ValueError("num_envs must be non-negative") + if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): + ids = list(range(num_envs)) + elif isinstance(env_ids, slice): + ids = list(range(num_envs))[env_ids] + elif isinstance(env_ids, torch.Tensor): + ids = env_ids.detach().cpu().reshape(-1).tolist() + else: + try: + ids = list(env_ids) + except TypeError: + ids = [env_ids] + normalized: list[int] = [] + for env_id in ids: + if isinstance(env_id, bool) or int(env_id) != env_id: + raise ValueError(f"Invalid environment ID: {env_id!r}") + env_id = int(env_id) + if env_id < 0 or env_id >= num_envs: + raise ValueError(f"Environment ID {env_id} is outside [0, {num_envs})") + normalized.append(env_id) + return normalized + + +def _select_texture_indices( + mode: str, + target_ids: list[int], + num_textures: int, + texture_indices: dict[int, int] | None, +) -> list[int]: + """Select one texture index for each target environment.""" + if num_textures <= 0: + raise ValueError("At least one texture source is required") + count = len(target_ids) + if mode == "random": + return torch.randint(num_textures, (count,)).tolist() + if mode == "without_replacement": + if count > num_textures: + raise ValueError( + "without_replacement requires at least one texture per target" + ) + return torch.randperm(num_textures)[:count].tolist() + if mode == "cycle": + return [index % num_textures for index in range(count)] + if mode == "fixed": + if texture_indices is None: + raise ValueError("fixed texture selection requires texture_indices") + selected: list[int] = [] + for env_id in target_ids: + if env_id not in texture_indices: + raise ValueError( + f"Missing fixed texture index for environment {env_id}" + ) + index = texture_indices[env_id] + if ( + isinstance(index, bool) + or int(index) != index + or not 0 <= int(index) < num_textures + ): + raise ValueError( + f"Invalid texture index for environment {env_id}: {index!r}" + ) + selected.append(int(index)) + return selected + raise ValueError(f"Unknown texture selection mode: {mode!r}") + + def set_rigid_object_visual_material( env: EmbodiedEnv, env_ids: Union[torch.Tensor, None], @@ -586,11 +673,12 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv): # Preload textures (currently only base color textures are supported) self.textures = [] + self.texture_refs = [] texture_path = get_data_path(cfg.params.get("texture_path", None)) if texture_path is not None: from embodichain.utils.utility import read_all_folder_images - texture_key = os.path.basename(texture_path) + texture_key = str(Path(texture_path).resolve()) # check if the texture group is already loaded in the global texture cache if texture_key in env.sim.get_texture_cache(): logger.log_info( @@ -615,6 +703,12 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv): env.sim.set_texture_cache(texture_key, self.textures) + # Create GPU texture handles once and reuse them on every reset. + # Passing image data to VisualMaterialInst would recreate a handle + # for each assignment, which is both expensive and invalidates the + # intended per-instance reuse semantics. + self.texture_refs = _get_texture_refs(env.sim, texture_key, self.textures) + if self.entity_cfg.uid == "default_plane": pass @@ -671,11 +765,17 @@ def gen_random_base_color_texture(width: int, height: int) -> torch.Tensor: rgba = (rgba * 255).to(torch.uint8) return rgba - def _randomize_texture(self, mat_inst: VisualMaterialInst) -> None: + def _randomize_texture( + self, mat_inst: VisualMaterialInst, texture_idx: int | None = None + ) -> 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 texture_idx is None: + texture_idx = torch.randint(0, len(self.textures), (1,)).item() + texture_refs = getattr(self, "texture_refs", None) + if texture_refs: + mat_inst.set_base_color_texture(texture_ref=texture_refs[texture_idx]) + else: + mat_inst.set_base_color_texture(texture_data=self.textures[texture_idx]) def _randomize_mat_inst( self, @@ -683,16 +783,17 @@ def _randomize_mat_inst( plan: Dict[str, torch.Tensor], random_texture_prob: float, idx: int = 0, + texture_idx: int | None = None, ) -> None: # randomize texture or base color based on the probability. - if random.random() < random_texture_prob and len(self.textures) != 0: - for key, value in plan.items(): - if key == "base_color": - mat_inst.set_base_color(value[idx].tolist()) - else: - getattr(mat_inst, f"set_{key}")(value[idx].item()) - - self._randomize_texture(mat_inst) + use_texture = random.random() < random_texture_prob and len(self.textures) != 0 + for key, value in plan.items(): + if key == "base_color": + mat_inst.set_base_color(value[idx].tolist()) + else: + getattr(mat_inst, f"set_{key}")(value[idx].item()) + if use_texture: + self._randomize_texture(mat_inst, texture_idx) else: # set a random base color instead. random_color_texture = ( @@ -711,19 +812,29 @@ def __call__( metallic_range: tuple[float, float] | None = None, roughness_range: tuple[float, float] | None = None, ior_range: tuple[float, float] | None = None, - ): + texture_sampling: str = "random", + texture_indices: Mapping[int, int] | None = None, + texture_scope: str = "per_material", + ) -> None: if self.entity_cfg.uid != "default_plane" and self.entity is None: return # resolve environment ids - if env_ids is None: - env_ids = torch.arange(env.num_envs, device="cpu") - else: - env_ids = env_ids.cpu() + env_ids = _normalize_env_ids(env_ids, env.num_envs) if self.entity_cfg.uid == "default_plane": env_ids = [0] + if texture_scope not in ("per_material", "per_instance"): + raise ValueError(f"Unsupported texture_scope: {texture_scope}") + texture_plan = ( + _select_texture_indices( + texture_sampling, env_ids, len(self.textures), texture_indices + ) + if self.textures + else [None] * len(env_ids) + ) + randomize_plan = {} if base_color_range: base_color = sample_uniform( @@ -767,10 +878,18 @@ def __call__( plan=randomize_plan, random_texture_prob=random_texture_prob, idx=0, + texture_idx=texture_plan[0] if texture_plan else None, ) return - for i, data in enumerate(self._mat_insts): + if isinstance(self.entity, RigidObject): + mat_insts = self.entity.get_visual_material_inst(env_ids=env_ids) + else: + 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): if isinstance(self.entity, RigidObject): # For RigidObject, data is the material instance directly mat: VisualMaterialInst = data @@ -784,6 +903,15 @@ def __call__( plan=randomize_plan, random_texture_prob=random_texture_prob, idx=i, + texture_idx=( + texture_plan[i] + if texture_plan + and ( + texture_scope == "per_instance" + or texture_sampling != "random" + ) + else None + ), ) else: for name, mat_inst in mat.items(): @@ -792,11 +920,17 @@ def __call__( plan=randomize_plan, random_texture_prob=random_texture_prob, idx=i, + texture_idx=( + texture_plan[i] + if texture_plan + and ( + texture_scope == "per_instance" + or texture_sampling != "random" + ) + else None + ), ) - env = self._env.sim.get_env() - env.clean_materials() - class randomize_indirect_lighting(Functor): """Randomize the environment's indirect (IBL) lighting or emissive light. diff --git a/embodichain/lab/sim/material.py b/embodichain/lab/sim/material.py index 7daddb8f2..afd17f73d 100644 --- a/embodichain/lab/sim/material.py +++ b/embodichain/lab/sim/material.py @@ -241,14 +241,32 @@ def set_emissive_intensity(self, intensity: float) -> None: logger.log_error("Unimplemented: set_emissive_intensity") def set_base_color_texture( - self, texture_path: str = None, texture_data: torch.Tensor | None = None + self, + texture_path: str = None, + texture_data: torch.Tensor | None = None, + texture_ref: object | None = None, ) -> None: - """Set base color texture from file path or texture data. + """Set base color texture from a file path, data, or existing texture. Args: texture_path: Path to texture file texture_data: Texture data as a torch.Tensor + texture_ref: Existing DexSim texture reference to assign directly. """ + if texture_ref is not None and ( + texture_path is not None or texture_data is not None + ): + logger.log_warning( + "texture_ref provided with texture_path or texture_data. " + "Using texture_ref." + ) + + if texture_ref is not None: + self.base_color_texture = texture_ref + inst = self._mat.get_inst(self.uid) + inst.set_base_color_map(texture_ref) + return + if texture_path is not None and texture_data is not None: logger.log_warning( "Both texture_path and texture_data are provided. Using texture_path." diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 89442b9a2..a0865155c 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -312,6 +312,7 @@ def __init__( # Global texture cache for material creation or randomization. # The structure is keys to the loaded texture data. The keys represent the texture group. self._texture_cache: Dict[str, Union[torch.Tensor, List[torch.Tensor]]] = dict() + self._texture_ref_cache: Dict[str, List[object]] = dict() self._init_sim_resources() @@ -802,6 +803,14 @@ def get_texture_cache( return None return self._texture_cache[key] + def get_texture_ref_cache(self, key: str) -> List[object] | None: + """Return cached GPU texture references for a canonical source key.""" + return self._texture_ref_cache.get(key) + + def set_texture_ref_cache(self, key: str, refs: List[object]) -> None: + """Cache GPU texture references for a canonical source key.""" + self._texture_ref_cache[key] = refs + def get_asset( self, uid: str ) -> Light | BaseSensor | Robot | RigidObject | Articulation | None: @@ -2769,6 +2778,7 @@ def _sever_wrapper_refs(obj_registry): self._visual_materials.clear() self._texture_cache.clear() + self._texture_ref_cache.clear() self._arenas.clear() self._markers.clear() self._gizmos.clear() diff --git a/tests/gym/envs/managers/test_randomize_visual_material.py b/tests/gym/envs/managers/test_randomize_visual_material.py new file mode 100644 index 000000000..06e69137c --- /dev/null +++ b/tests/gym/envs/managers/test_randomize_visual_material.py @@ -0,0 +1,220 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest +import torch +import embodichain.lab.gym.envs.managers.randomization.visual as visual +import dexsim + +from embodichain.lab.sim.material import VisualMaterialInst +from embodichain.lab.gym.envs.managers.randomization.visual import ( + _normalize_env_ids, + _select_texture_indices, + randomize_visual_material, +) + + +class FakeMaterialInstance: + def __init__(self): + self.base_color_map = None + + def set_base_color_map(self, texture): + self.base_color_map = texture + + +class FakeMaterial: + def __init__(self): + self.instances = {} + + def get_inst(self, uid): + return self.instances.setdefault(uid, FakeMaterialInstance()) + + +def test_set_base_color_texture_uses_texture_ref(): + material = FakeMaterial() + instance = VisualMaterialInst("instance", material) + texture_ref = object() + + instance.set_base_color_texture(texture_ref=texture_ref) + + assert material.get_inst("instance").base_color_map is texture_ref + assert instance.base_color_texture is texture_ref + + +def test_normalize_env_ids_supports_all_input_forms(): + assert _normalize_env_ids(None, 3) == [0, 1, 2] + assert _normalize_env_ids(torch.tensor([2, 0]), 3) == [2, 0] + assert _normalize_env_ids([1], 3) == [1] + assert _normalize_env_ids(slice(None), 3) == [0, 1, 2] + + +def test_texture_selection_modes(): + assert sorted( + _select_texture_indices("without_replacement", [0, 1, 2], 3, None) + ) == [0, 1, 2] + assert _select_texture_indices("fixed", [3, 1], 4, {1: 2, 3: 0}) == [0, 2] + with pytest.raises(ValueError, match="without_replacement"): + _select_texture_indices("without_replacement", [0, 1], 1, None) + + +def test_partial_reset_targets_only_selected_environment_ids(monkeypatch): + class Obj: + is_shared_visual_material = False + + def __init__(self): + self.mats = [FakeMat() for _ in range(4)] + + def get_visual_material_inst(self, env_ids=None, **kwargs): + return [self.mats[i] for i in env_ids] + + functor = object.__new__(randomize_visual_material) + functor.entity_cfg = type("C", (), {"uid": "x", "link_names": None})() + functor.entity = Obj() + functor.textures = [] + monkeypatch.setattr(visual, "RigidObject", Obj) + + def mark(*, mat_inst, **kwargs): + mat_inst.set_base_color([1, 1, 1, 1]) + + functor._randomize_mat_inst = mark + env = type("E", (), {"num_envs": 4})() + functor.__call__( + env, torch.tensor([1, 3]), functor.entity_cfg, random_texture_prob=0 + ) + assert functor.entity.mats[0].color is None and functor.entity.mats[2].color is None + assert ( + functor.entity.mats[1].color is not None + and functor.entity.mats[3].color is not None + ) + + +def test_fixed_assignment_maps_global_environment_ids(monkeypatch): + class Obj: + is_shared_visual_material = False + + def get_visual_material_inst(self, env_ids=None, **kwargs): + return [{"l": FakeMat()} for _ in env_ids] + + f = object.__new__(randomize_visual_material) + f.entity_cfg = type("C", (), {"uid": "x", "link_names": None})() + f.entity = Obj() + f.textures = ["a", "b", "c"] + monkeypatch.setattr(visual, "RigidObject", Obj) + seen = [] + + def record(*, texture_idx, **kwargs): + seen.append(texture_idx) + + f._randomize_mat_inst = record + f.__call__( + type("E", (), {"num_envs": 4})(), + torch.tensor([1, 3]), + f.entity_cfg, + texture_sampling="fixed", + texture_indices={1: 2, 3: 0}, + ) + assert seen == [2, 0] + + +def test_per_instance_selection_is_reused_for_all_links(monkeypatch): + class Obj: + is_shared_visual_material = False + + def get_visual_material_inst(self, env_ids=None, **kwargs): + return [{"a": FakeMat(), "b": FakeMat()} for _ in env_ids] + + f = object.__new__(randomize_visual_material) + f.entity_cfg = type("C", (), {"uid": "x", "link_names": None})() + f.entity = Obj() + f.textures = ["a", "b"] + monkeypatch.setattr(visual, "RigidObject", type("Other", (), {})) + monkeypatch.setattr(visual, "Articulation", Obj) + seen = [] + + def record(*, texture_idx, **kwargs): + seen.append(texture_idx) + + f._randomize_mat_inst = record + f.__call__( + type("E", (), {"num_envs": 4})(), + torch.tensor([1, 3]), + f.entity_cfg, + texture_sampling="fixed", + texture_indices={1: 1, 3: 0}, + texture_scope="per_instance", + ) + assert seen == [1, 1, 0, 0] + + +class FakeMat: + def __init__(self): + self.color = None + + def set_base_color(self, value): + self.color = value + + +def test_generated_color_branch_sets_texture_without_unbound_error(): + functor = object.__new__(randomize_visual_material) + texture = object() + functor.textures = [texture] + + class Instance: + texture = None + + def set_base_color_texture(self, texture_data=None, **kwargs): + self.texture = texture_data + + instance = Instance() + functor._randomize_mat_inst(instance, {}, random_texture_prob=1.0) + assert instance.texture is texture + + +def test_texture_references_are_created_once(monkeypatch): + calls = [] + + class Sim: + def __init__(self): + self.cache = {} + + def get_texture_ref_cache(self, key): + return self.cache.get(key) + + def set_texture_ref_cache(self, key, refs): + self.cache[key] = refs + + class E: + def create_color_texture(self, image, has_alpha): + calls.append(image) + return object() + + monkeypatch.setattr( + dexsim, "default_world", lambda: type("W", (), {"get_env": lambda s: E()})() + ) + sim = Sim() + images = [torch.zeros((2, 2, 4), dtype=torch.uint8)] + visual._get_texture_refs(sim, "/tmp/source", images) + visual._get_texture_refs(sim, "/tmp/source", images) + assert len(calls) == 1 + + +def test_four_environment_fake_assignments_cover_reset_cases(): + first = _select_texture_indices("without_replacement", [0, 1, 2, 3], 4, None) + assert len(set(first)) == 4 + assert _select_texture_indices("fixed", [1, 3], 4, {1: 2, 3: 0}) == [2, 0] + assert _select_texture_indices("fixed", [3], 4, {3: 1}) == [1]