diff --git a/docs/source/overview/gym/event_functors.md b/docs/source/overview/gym/event_functors.md index 46ed991a..fb110ef3 100644 --- a/docs/source/overview/gym/event_functors.md +++ b/docs/source/overview/gym/event_functors.md @@ -83,6 +83,27 @@ This page lists all available event functors that can be used with the Event Man "params": {"color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], "intensity_range": [0.5, 2.0]}} ``` +* - {class}`~randomization.visual.randomize_indirect_lighting` + - Randomize indirect (IBL) lighting or emissive light. Implemented as a Functor class. Operates in one of two **mutually exclusive** modes — configuring both raises a ``ValueError``: + + **HDR mode** — provide ``path`` pointing to a folder of ``.hdr`` files. A random file is selected on each call and applied as the environment map. The ``path`` is resolved via ``get_data_path``, supporting absolute paths, data-root-relative paths, and dataset-class paths. + + ```json + {"func": "randomize_indirect_lighting", + "mode": "interval", "interval_step": 10, + "params": {"path": "EnvMapHDR/EnvMapHDR"}} + ``` + + **Emissive mode** — provide ``emissive_color_range`` (pair of RGB lists) and/or ``emissive_intensity_range`` (pair of floats). Color and intensity are sampled uniformly on each call and applied via ``set_emission_light``. + + ```json + {"func": "randomize_indirect_lighting", + "mode": "interval", "interval_step": 10, + "params": {"emissive_color_range": [[0.8, 0.8, 0.8], [1.0, 1.0, 1.0]], + "emissive_intensity_range": [80.0, 150.0]}} + ``` + + Applies the same lighting to all environments. * - {func}`~randomization.visual.randomize_camera_extrinsics` - Randomize camera poses for viewpoint diversity. Supports both attach mode (pos/euler perturbation) and look_at mode (eye/target/up perturbation). diff --git a/embodichain/data/assets/materials.py b/embodichain/data/assets/materials.py index ced7f82a..22183147 100644 --- a/embodichain/data/assets/materials.py +++ b/embodichain/data/assets/materials.py @@ -100,6 +100,46 @@ def get_material_list(self) -> List[str]: ] +class EnvMapHDR(EmbodiChainDataset): + def __init__(self, data_root: str = None): + data_descriptor = o3d.data.DataDescriptor( + os.path.join(EMBODICHAIN_DOWNLOAD_PREFIX, material_assets, "EnvMapHDR.zip"), + "ea7abc8e955fe64069073d63834da60e", + ) + prefix = type(self).__name__ + path = EMBODICHAIN_DEFAULT_DATA_ROOT if data_root is None else data_root + + super().__init__(prefix, data_descriptor, path) + + def get_env_map_path(self, name: str) -> str: + """Get the path of an HDR environment map. + + Args: + name (str): The name of the HDR environment map. + + Returns: + str: The path to the HDR environment map file. + """ + env_map_names = self.get_env_map_list() + if name not in env_map_names: + logger.log_error( + f"Invalid env map name: {name}. Available names are: {env_map_names}" + ) + return str(Path(self.extract_dir) / "EnvMapHDR" / name) + + def get_env_map_list(self) -> List[str]: + """Get the names of all HDR environment maps. + + Returns: + List[str]: The names of all HDR environment map files. + """ + return [ + f.name + for f in Path(self.extract_dir).glob("EnvMapHDR/*.hdr") + if f.is_file() + ] + + class CocoBackground(EmbodiChainDataset): def __init__(self, data_root: str = None): data_descriptor = o3d.data.DataDescriptor( diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index d6ca36d9..18137d87 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -54,6 +54,7 @@ init_rollout_buffer_from_gym_space, ) from embodichain.utils import configclass, logger +from embodichain.data import get_data_path __all__ = ["EmbodiedEnvCfg", "EmbodiedEnv"] @@ -932,6 +933,9 @@ def _setup_lights(self) -> None: if self.cfg.light.indirect is not None: if "emission_light" in self.cfg.light.indirect: self.sim.set_emission_light(**self.cfg.light.indirect["emission_light"]) + if "env_map" in self.cfg.light.indirect: + path = get_data_path(self.cfg.light.indirect["env_map"]) + self.sim.set_indirect_lighting(path) def _setup_background(self) -> None: """Setup the static rigid objects in the environment.""" diff --git a/embodichain/lab/gym/envs/managers/randomization/visual.py b/embodichain/lab/gym/envs/managers/randomization/visual.py index 17daa5d4..c49707bd 100644 --- a/embodichain/lab/gym/envs/managers/randomization/visual.py +++ b/embodichain/lab/gym/envs/managers/randomization/visual.py @@ -21,6 +21,7 @@ import random import copy import numpy as np +from pathlib import Path from typing import TYPE_CHECKING, Literal, Union, Dict @@ -59,6 +60,7 @@ "set_rigid_object_visual_material", "set_rigid_object_group_visual_material", "randomize_visual_material", + "randomize_indirect_lighting", ] @@ -742,3 +744,134 @@ def __call__( env = self._env.sim.get_env() env.clean_materials() + + +class randomize_indirect_lighting(Functor): + """Randomize the environment's indirect (IBL) lighting or emissive light. + + This functor operates in one of two mutually exclusive modes: + + * **HDR mode** — ``path`` is provided. A random ``.hdr`` file is chosen from + the folder on every call and applied via :meth:`set_indirect_lighting`. + * **Emissive mode** — ``emissive_color_range`` and/or + ``emissive_intensity_range`` are provided. The emissive light color and + intensity are sampled uniformly on every call and applied via + :meth:`set_emission_light`. + + Providing both ``path`` and emissive parameters simultaneously is an error. + + .. attention:: + This functor applies the same lighting to all environments. + + .. tip:: + The ``path`` parameter is resolved via :func:`get_data_path`, so it + supports absolute paths, data-root-relative paths, and dataset-class + paths (e.g. ``"EnvMapHDR"``). + + ``emissive_color_range`` is a pair of ``[r, g, b]`` lists representing + the lower and upper bounds for sampling the emissive color, e.g. + ``[[0.8, 0.8, 0.8], [1.0, 1.0, 1.0]]``. + + ``emissive_intensity_range`` is a ``[min, max]`` pair for the emissive + intensity scalar, e.g. ``[80.0, 150.0]``. + """ + + def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv): + """Initialize the functor. + + Args: + cfg: The configuration of the functor. + + * **HDR mode**: set ``params["path"]`` to a folder of ``.hdr`` files. + * **Emissive mode**: set ``params["emissive_color_range"]`` + (pair of RGB lists) and/or ``params["emissive_intensity_range"]`` + (pair of floats). + + env: The environment instance. + + Raises: + ValueError: If both HDR and emissive params are provided, or if + neither is provided. + """ + super().__init__(cfg, env) + + has_hdr = cfg.params.get("path", None) is not None + has_emissive = ( + cfg.params.get("emissive_color_range", None) is not None + or cfg.params.get("emissive_intensity_range", None) is not None + ) + + if has_hdr and has_emissive: + raise ValueError( + "randomize_indirect_lighting: 'path' (HDR mode) and emissive " + "parameters ('emissive_color_range', 'emissive_intensity_range') " + "are mutually exclusive. Configure only one mode." + ) + if not has_hdr and not has_emissive: + raise ValueError( + "randomize_indirect_lighting: provide either 'path' for HDR " + "mode, or 'emissive_color_range'/'emissive_intensity_range' for " + "emissive mode." + ) + + # HDR mode state + self._hdr_files: list[Path] = [] + if has_hdr: + path = get_data_path(cfg.params["path"]) + self._hdr_files = sorted(Path(path).glob("*.hdr")) + if not self._hdr_files: + logger.log_warning( + f"No .hdr files found in '{path}'. " + f"Indirect lighting randomization will be a no-op." + ) + + # Emissive mode state + self._emissive_color_range: tuple[list[float], list[float]] | None = ( + cfg.params.get("emissive_color_range", None) + ) + self._emissive_intensity_range: tuple[float, float] | None = cfg.params.get( + "emissive_intensity_range", None + ) + + def __call__( + self, + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, None], + path: str | None = None, + ) -> None: + """Randomize lighting according to the configured mode. + + In HDR mode a random ``.hdr`` file is selected and applied. In emissive + mode the emissive color and/or intensity are sampled and applied. + + Args: + env: The environment instance. + env_ids: Target environment IDs (unused — lighting is global). + path: Ignored. Kept for interface compatibility with the event system. + """ + if self._hdr_files: + # HDR mode + selected = random.choice(self._hdr_files) + env.sim.set_indirect_lighting(str(selected)) + return + + # Emissive mode + emissive_color: list[float] | None = None + if self._emissive_color_range is not None: + color_tensor = sample_uniform( + lower=torch.tensor(self._emissive_color_range[0]), + upper=torch.tensor(self._emissive_color_range[1]), + size=(1, 3), + ) + emissive_color = color_tensor.squeeze(0).tolist() + + emissive_intensity: float | None = None + if self._emissive_intensity_range is not None: + emissive_intensity = float( + np.random.uniform( + self._emissive_intensity_range[0], + self._emissive_intensity_range[1], + ) + ) + + env.sim.set_emission_light(color=emissive_color, intensity=emissive_intensity) diff --git a/embodichain/lab/scripts/preview_asset.py b/embodichain/lab/scripts/preview_asset.py index bef02faa..49c86de5 100644 --- a/embodichain/lab/scripts/preview_asset.py +++ b/embodichain/lab/scripts/preview_asset.py @@ -34,6 +34,16 @@ python -m embodichain.lab.scripts.preview_asset \\ --asset_path /path/to/asset.usda \\ --headless + + # Preview with a built-in environment map + python -m embodichain.lab.scripts.preview_asset \\ + --asset_path /path/to/sugar_box.usda \\ + --env_map "Studio" + + # Preview with a custom HDR environment map + python -m embodichain.lab.scripts.preview_asset \\ + --asset_path /path/to/sugar_box.usda \\ + --env_map /path/to/environment.hdr """ from __future__ import annotations @@ -208,6 +218,10 @@ def main(args: argparse.Namespace) -> None: sim = SimulationManager(sim_cfg) try: + if args.env_map: + log_info(f"Setting environment map: {args.env_map} ...", color="green") + sim.set_indirect_lighting(args.env_map) + assets = load_assets(sim, args) log_info(f"Loaded {len(assets)} asset(s) successfully.", color="green") @@ -318,6 +332,15 @@ def cli(): default="hybrid", help="Renderer backend (default: hybrid).", ) + parser.add_argument( + "--env_map", + type=str, + default=None, + help=( + "Environment map for indirect lighting. Accepts a built-in IBL resource " + "name (e.g. 'Studio') or an absolute file path (.hdr/.png/.exr)." + ), + ) parser.add_argument( "--preview", action="store_true", diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 6f1ba901..8f2e257f 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -658,6 +658,17 @@ def set_default_background(self) -> None: self._default_plane.set_material(mat.get_instance("plane_mat").mat) self._visual_materials[mat_name] = mat + def set_ground_plane_visibility(self, visible: bool) -> None: + """_summary_ + + Args: + visible (bool): _description_ + """ + if visible: + self._default_plane.set_visible(True) + else: + self._default_plane.set_visible(False) + def set_texture_cache( self, key: str, texture: Union[torch.Tensor, List[torch.Tensor]] ) -> None: diff --git a/examples/sim/scene/scene_demo.py b/examples/sim/scene/scene_demo.py index b119cdfb..1c08af6a 100644 --- a/examples/sim/scene/scene_demo.py +++ b/examples/sim/scene/scene_demo.py @@ -126,7 +126,7 @@ def main(): num_lights = 8 radius = 5 height = 8 - intensity = 200 + intensity = 50 lights = [] for i in range(num_lights): diff --git a/tests/gym/envs/managers/test_event_functors.py b/tests/gym/envs/managers/test_event_functors.py index e7e206de..981e44ab 100644 --- a/tests/gym/envs/managers/test_event_functors.py +++ b/tests/gym/envs/managers/test_event_functors.py @@ -283,6 +283,13 @@ def get_rigid_object_group(self, uid: str): def update(self, step: int = 1): pass + def set_indirect_lighting(self, path: str) -> None: + self._last_indirect_lighting = path + + def set_emission_light(self, color=None, intensity=None) -> None: + self._last_emission_color = color + self._last_emission_intensity = intensity + class MockEnv: """Mock environment for event functor tests.""" @@ -324,6 +331,10 @@ def __init__(self, num_envs: int = 4, num_joints: int = 6): from embodichain.lab.gym.envs.managers.randomization.spatial import ( randomize_articulation_root_pose, ) +from embodichain.lab.gym.envs.managers.randomization.visual import ( + randomize_indirect_lighting, +) +from embodichain.lab.gym.envs.managers import FunctorCfg class TestResolveUids: @@ -815,3 +826,187 @@ def test_handles_nonexistent_link_pattern(self): mass_range=(0.5, 2.0), link_names="nonexistent_link", ) + + +class TestRandomizeIndirectLighting: + """Tests for the randomize_indirect_lighting functor.""" + + def _make_cfg(self, params: dict) -> FunctorCfg: + cfg = FunctorCfg(func=randomize_indirect_lighting) + cfg.params = params + return cfg + + # ------------------------------------------------------------------ + # Init validation + # ------------------------------------------------------------------ + + def test_raises_when_no_params(self, tmp_path): + """Raises ValueError when neither HDR path nor emissive params given.""" + env = MockEnv() + cfg = self._make_cfg({}) + with pytest.raises(ValueError, match="provide either"): + randomize_indirect_lighting(cfg, env) + + def test_raises_when_both_hdr_and_emissive(self, tmp_path): + """Raises ValueError when HDR path and emissive params are both set.""" + hdr_dir = tmp_path / "hdr" + hdr_dir.mkdir() + (hdr_dir / "a.hdr").write_text("") + env = MockEnv() + cfg = self._make_cfg( + { + "path": str(hdr_dir), + "emissive_color_range": [[0.5, 0.5, 0.5], [1.0, 1.0, 1.0]], + } + ) + with pytest.raises(ValueError, match="mutually exclusive"): + randomize_indirect_lighting(cfg, env) + + # ------------------------------------------------------------------ + # HDR mode + # ------------------------------------------------------------------ + + def test_hdr_mode_calls_set_indirect_lighting(self, tmp_path): + """HDR mode calls set_indirect_lighting with one of the .hdr files.""" + hdr_dir = tmp_path / "hdr" + hdr_dir.mkdir() + files = ["sky1.hdr", "sky2.hdr", "sky3.hdr"] + for f in files: + (hdr_dir / f).write_text("") + env = MockEnv() + cfg = self._make_cfg({"path": str(hdr_dir)}) + functor = randomize_indirect_lighting(cfg, env) + + functor(env, None) + + chosen = env.sim._last_indirect_lighting + assert chosen.endswith(".hdr") + assert any(chosen.endswith(f) for f in files) + + def test_hdr_mode_does_not_call_set_emission_light(self, tmp_path): + """HDR mode must not touch emissive light.""" + hdr_dir = tmp_path / "hdr" + hdr_dir.mkdir() + (hdr_dir / "sky.hdr").write_text("") + env = MockEnv() + cfg = self._make_cfg({"path": str(hdr_dir)}) + functor = randomize_indirect_lighting(cfg, env) + + # Ensure attribute not set by HDR call + env.sim._last_emission_color = "sentinel" + env.sim._last_emission_intensity = "sentinel" + + functor(env, None) + + assert env.sim._last_emission_color == "sentinel" + assert env.sim._last_emission_intensity == "sentinel" + + def test_hdr_mode_noop_when_no_hdr_files(self, tmp_path): + """HDR mode is a no-op (no crash) when the folder has no .hdr files.""" + hdr_dir = tmp_path / "empty" + hdr_dir.mkdir() + env = MockEnv() + cfg = self._make_cfg({"path": str(hdr_dir)}) + functor = randomize_indirect_lighting(cfg, env) + + functor(env, None) # must not raise + + assert not hasattr(env.sim, "_last_indirect_lighting") + + def test_hdr_mode_selects_from_available_files(self, tmp_path): + """HDR mode always selects a file from the provided folder over many calls.""" + hdr_dir = tmp_path / "hdr" + hdr_dir.mkdir() + names = [f"env{i}.hdr" for i in range(5)] + for n in names: + (hdr_dir / n).write_text("") + env = MockEnv() + cfg = self._make_cfg({"path": str(hdr_dir)}) + functor = randomize_indirect_lighting(cfg, env) + + chosen_set = set() + for _ in range(50): + functor(env, None) + chosen_set.add(env.sim._last_indirect_lighting) + + # All chosen paths must be valid HDR files from the folder + valid_paths = {str(hdr_dir / n) for n in names} + assert chosen_set.issubset(valid_paths) + + # ------------------------------------------------------------------ + # Emissive mode + # ------------------------------------------------------------------ + + def test_emissive_color_mode_calls_set_emission_light(self): + """Emissive mode calls set_emission_light with color in range.""" + env = MockEnv() + cfg = self._make_cfg( + {"emissive_color_range": [[0.2, 0.3, 0.4], [0.6, 0.7, 0.8]]} + ) + functor = randomize_indirect_lighting(cfg, env) + + functor(env, None) + + color = env.sim._last_emission_color + assert color is not None + assert len(color) == 3 + assert 0.2 <= color[0] <= 0.6 + assert 0.3 <= color[1] <= 0.7 + assert 0.4 <= color[2] <= 0.8 + assert env.sim._last_emission_intensity is None + + def test_emissive_intensity_mode_calls_set_emission_light(self): + """Emissive mode calls set_emission_light with intensity in range.""" + env = MockEnv() + cfg = self._make_cfg({"emissive_intensity_range": [50.0, 150.0]}) + functor = randomize_indirect_lighting(cfg, env) + + functor(env, None) + + assert env.sim._last_emission_color is None + intensity = env.sim._last_emission_intensity + assert intensity is not None + assert 50.0 <= intensity <= 150.0 + + def test_emissive_color_and_intensity_together(self): + """Both color and intensity can be set together in emissive mode.""" + env = MockEnv() + cfg = self._make_cfg( + { + "emissive_color_range": [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]], + "emissive_intensity_range": [80.0, 120.0], + } + ) + functor = randomize_indirect_lighting(cfg, env) + + functor(env, None) + + color = env.sim._last_emission_color + intensity = env.sim._last_emission_intensity + assert color is not None and len(color) == 3 + assert all(0.0 <= c <= 1.0 for c in color) + assert 80.0 <= intensity <= 120.0 + + def test_emissive_mode_does_not_call_set_indirect_lighting(self): + """Emissive mode must not touch indirect lighting (HDR).""" + env = MockEnv() + cfg = self._make_cfg({"emissive_intensity_range": [100.0, 100.0]}) + functor = randomize_indirect_lighting(cfg, env) + + functor(env, None) + + assert not hasattr(env.sim, "_last_indirect_lighting") + + def test_emissive_values_vary_across_calls(self): + """Emissive intensity is sampled fresh on each call (not fixed).""" + env = MockEnv() + cfg = self._make_cfg({"emissive_intensity_range": [0.0, 1000.0]}) + functor = randomize_indirect_lighting(cfg, env) + + intensities = set() + for _ in range(20): + functor(env, None) + intensities.add(round(env.sim._last_emission_intensity, 4)) + + # Over 20 draws from [0, 1000] all values being identical is astronomically unlikely + assert len(intensities) > 1