diff --git a/agent_context/topics/randomization/randomization.md b/agent_context/topics/randomization/randomization.md index b7f894a65..057206007 100644 --- a/agent_context/topics/randomization/randomization.md +++ b/agent_context/topics/randomization/randomization.md @@ -42,13 +42,16 @@ The `__init__.py` of the randomization package re-exports everything via `from . | `randomize_visual_material` | Random material properties | (varies) | | `randomize_camera_extrinsics` | Camera pose | `pos_range`, `euler_range` (attach mode) or `eye_range`, `target_range`, `up_range` (look-at mode) | | `randomize_camera_intrinsics` | Camera intrinsic params | (varies) | -| `randomize_light` | Light pos/color/intensity | `position_range`, `color_range`, `intensity_range` | +| `randomize_light` | Light pos/color/intensity/direction | `position_range`, `color_range`, `intensity_range`, `direction_range` | | `randomize_emission_light` | Emission light props | (varies) | | `randomize_indirect_lighting` | Indirect lighting | (varies) | - Camera extrinsics auto-detect mode: if `extrinsics.parent` is set → attach mode (pos/euler perturbation via `set_local_pose`); if `extrinsics.eye` is set → look-at mode (eye/target/up perturbation via `look_at`). - `set_rigid_object_visual_material` is deterministic (not random) but uses the same functor mechanism for fixed material assignment at reset. - Light randomization applies the **same values across all envs** (documented limitation). +- ``position_range`` is ignored for global scene lights (``"sun"``, ``"direction"``). Use ``direction_range`` instead. +- ``direction_range`` is only applicable for directional light types (``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, ``"mesh"``). +- Global lights (``"sun"``, ``"direction"``) have ``num_instances == 1``; all other types are batched per environment. ### Spatial (`spatial.py`) diff --git a/docs/source/overview/gym/event_functors.md b/docs/source/overview/gym/event_functors.md index eab6ad954..2ab03ddfc 100644 --- a/docs/source/overview/gym/event_functors.md +++ b/docs/source/overview/gym/event_functors.md @@ -64,7 +64,15 @@ This page lists all available event functors that can be used with the Event Man "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]]}} ``` * - {func}`~randomization.visual.randomize_light` - - Vary light position, color, and intensity within specified ranges. + - Vary light position, color, intensity, and direction within specified ranges. + + .. note:: + ``position_range`` is ignored for global scene lights (``"sun"``, ``"direction"``). + Use ``direction_range`` instead for these types. + + .. note:: + ``direction_range`` is only applicable for directional light types + (``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, ``"mesh"``). ```json {"func": "randomize_light", @@ -72,7 +80,8 @@ This page lists all available event functors that can be used with the Event Man "params": {"entity_cfg": {"uid": "light_1"}, "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], - "intensity_range": [50.0, 100.0]}} + "intensity_range": [50.0, 100.0], + "direction_range": [[-0.1, -0.1, -0.1], [0.1, 0.1, 0.1]]}} ``` * - {func}`~randomization.visual.randomize_emission_light` - Randomize global emission light color and intensity. Applies the same emission light properties across all environments. diff --git a/docs/source/overview/sim/sim_assets.md b/docs/source/overview/sim/sim_assets.md index ab49dfbbb..fca7351f4 100644 --- a/docs/source/overview/sim/sim_assets.md +++ b/docs/source/overview/sim/sim_assets.md @@ -112,14 +112,27 @@ For Rigid Object tutorial, please refer to the [Create Scene](https://dexforce.g ### Lights -Configured via `LightCfg`. +Configured via `LightCfg`. Supports six light types matching the dexsim rendering backend. | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | -| `light_type` | `Literal` | `"point"` | Type of light (currently only "point"). | -| `color` | `tuple` | `(1.0, 1.0, 1.0)` | RGB color. | -| `intensity` | `float` | `50.0` | Intensity in watts/m^2. | -| `radius` | `float` | `1e2` | Falloff radius. | +| `light_type` | `Literal` | `"point"` | Light type: ``"point"``, ``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, or ``"mesh"``. | +| `color` | `tuple` | `(1.0, 1.0, 1.0)` | RGB color of the light source. | +| `intensity` | `float` | `30.0` | Intensity of the light source in watts/m^2. | +| `enable_shadow` | `bool` | `True` | Whether the light casts shadows. | +| `radius` | `float` | `10.0` | Falloff radius (only for ``"point"``). | +| `direction` | `tuple` | `(0.0, 0.0, -1.0)` | Direction vector for directional types (``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, ``"mesh"``). | +| `spot_angle_inner` | `float` | `30.0` | Inner cone angle in degrees (only for ``"spot"``). | +| `spot_angle_outer` | `float` | `45.0` | Outer cone angle in degrees (only for ``"spot"``). | +| `rect_width` | `float` | `1.0` | Width of rectangular area light (only for ``"rect"``). | +| `rect_height` | `float` | `1.0` | Height of rectangular area light (only for ``"rect"``). | +| `mesh_path` | `str` | `""` | Asset path for mesh-based emissive lights (only for ``"mesh"``). | + +.. attention:: + ``"sun"`` and ``"direction"`` are **global scene lights** (infinite-distance directional light + sources). They are created as a single instance on the root environment, not batched per + environment. Use :meth:`Light.set_direction` instead of :meth:`Light.set_local_pose` for + these types. ```{toctree} diff --git a/docs/source/overview/sim/sim_manager.md b/docs/source/overview/sim/sim_manager.md index 675f8a06d..9a4a50c30 100644 --- a/docs/source/overview/sim/sim_manager.md +++ b/docs/source/overview/sim/sim_manager.md @@ -103,6 +103,79 @@ sim_config = SimulationManagerCfg( ) ``` +### DLSS 3.5 (OfflineRT Only) + +When `render_cfg.renderer="rt"` (OfflineRT) and the simulation is running with a visible window, you can enable NVIDIA DLSS 3.5 for AI-powered denoising and upscaling through `render_cfg.dlss`. + +DLSS 3.5 combines two features: + +- **Ray Reconstruction (RR)** — replaces the OptiX denoiser for the window camera. +- **Super Resolution (SR)** — upscales the RR output from render resolution to the target/display resolution. + +Both features are enabled together when `dlss_enabled=True`. The exact upscale ratio can be controlled either with the `dlss_quality` preset or with the `upsample_ratio` convenience parameter. + +| Parameter | Type | Default | Description | +| :--- | :--- | :--- | :--- | +| `dlss_enabled` | `bool` | `True` | Master DLSS enable toggle. | +| `dlss_quality` | `int` | `-1` | DLSS quality preset. `-1`=auto, `0`=UltraPerformance (~3.0×), `1`=Performance (~2.25×), `2`=Balanced (~1.75×), `3`=Quality (~1.5×), `4`=UltraQuality (~1.3×), `5`=DLAA (1.0×). | +| `upsample_ratio` | `float` | `1.0` | Convenience ratio `target_width / render_width` (and height). When `render_width`/`render_height` are `0`, the render resolution is computed as `target / upsample_ratio`. Defaults to `1.0` (1:1 / DLAA mode). Must be `>= 1.0`. | +| `render_width` | `int` | `0` | Internal render resolution width. `0` means use the effective target width. | +| `render_height` | `int` | `0` | Internal render resolution height. `0` means use the effective target height. | +| `target_width` | `int` | `0` | Target/display resolution width. `0` means use the window width. | +| `target_height` | `int` | `0` | Target/display resolution height. `0` means use the window height. | +| `exposure_compensation` | `float` | `1.0` | Multiplier on the DLSS-RR pre-exposure scalar. `>1` brightens the HDR input to RR, which helps reject ghosting in dark areas; `<1` darkens. Typical range `0.5–8.0`. | + +#### Resolution selection rules + +- The **target** resolution defaults to the window size unless `target_width`/`target_height` are set explicitly. +- The **render** resolution defaults to the target resolution (1:1 / DLAA mode) unless overridden. +- If `upsample_ratio > 1.0` and `render_width`/`render_height` are `0`, the render resolution is computed as `target / upsample_ratio`. At `1.0`, render defaults to the target resolution (1:1 / DLAA mode). +- Explicit `render_width`/`render_height` or `target_width`/`target_height` take precedence over computed values. +- The window is always resized to match the effective target resolution when DLSS is active. + +#### Examples + +Use the `upsample_ratio` shortcut to render at half resolution and upscale to the window size: + +```python +from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim.cfg import DLSSCfg, RenderCfg + +sim_config = SimulationManagerCfg( + width=1920, + height=1080, + render_cfg=RenderCfg( + renderer="rt", + dlss=DLSSCfg( + dlss_enabled=True, + upsample_ratio=2.0, # render at 960x540, upscale to 1920x1080 + ), + ), +) +``` + +Set explicit render and target resolutions: + +```python +from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim.cfg import DLSSCfg, RenderCfg + +sim_config = SimulationManagerCfg( + render_cfg=RenderCfg( + renderer="rt", + dlss=DLSSCfg( + dlss_enabled=True, + render_width=1280, + render_height=720, + target_width=1920, + target_height=1080, + ), + ), +) +``` + +> **Note:** Offscreen cameras always use the OptiX denoiser regardless of DLSS settings. DLSS only affects the interactive viewer window when the OfflineRT renderer is used. + ## Initialization diff --git a/docs/source/tutorial/modular_env.rst b/docs/source/tutorial/modular_env.rst index eef801c39..af4d14992 100644 --- a/docs/source/tutorial/modular_env.rst +++ b/docs/source/tutorial/modular_env.rst @@ -129,11 +129,12 @@ Configures a stereo camera system using :class:`StereoCameraCfg`: :language: python :lines: 120-130 -Defines scene illumination with controllable point lights: +Defines scene illumination with configurable lights: -- **Type**: Point light for realistic shadows -- **Properties**: Configurable color, intensity, and position -- **UID**: Named reference for event system manipulation +- **Types**: Supports ``"point"``, ``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, and ``"mesh"``. +- **Global lights**: ``"sun"`` and ``"direction"`` are global scene lights (single instance, infinite distance). +- **Properties**: Configurable color, intensity, position, and direction. +- **UID**: Named reference for event system manipulation. **Rigid Objects** diff --git a/embodichain/lab/gym/envs/managers/randomization/visual.py b/embodichain/lab/gym/envs/managers/randomization/visual.py index c49707bd5..bb4a08ac9 100644 --- a/embodichain/lab/gym/envs/managers/randomization/visual.py +++ b/embodichain/lab/gym/envs/managers/randomization/visual.py @@ -272,6 +272,7 @@ def randomize_light( position_range: tuple[list[float], list[float]] | None = None, color_range: tuple[list[float], list[float]] | None = None, intensity_range: tuple[float, float] | None = None, + direction_range: tuple[list[float], list[float]] | None = None, ) -> None: """Randomize light properties by adding, scaling, or setting random values. @@ -289,10 +290,17 @@ def randomize_light( position_range is the x, y, z value added into light's cfg.init_pos. color_range is the absolute r, g, b value set to the light object. intensity_range is the value added into light's cfg.intensity. + direction_range is the x, y, z value added into light's cfg.direction. + (Only applicable for ``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, and ``"mesh"`` light types.) .. tip:: This function uses CPU tensors to assign light properties. + .. warning:: + ``position_range`` is ignored for global scene lights (``"sun"``, ``"direction"``) + because they are infinite-distance lights with no meaningful position. + Use ``direction_range`` instead for these light types. + Args: env (EmbodiedEnv): The environment instance. env_ids (Union[torch.Tensor, None]): The environment IDs to apply the randomization. @@ -300,25 +308,45 @@ def randomize_light( position_range (tuple[list[float], list[float]] | None): The range for the position randomization. color_range (tuple[list[float], list[float]] | None): The range for the color randomization. intensity_range (tuple[float, float] | None): The range for the intensity randomization. + direction_range (tuple[list[float], list[float]] | None): The range for the direction randomization. + Only applicable for directional light types (``"sun"``, ``"direction"``, ``"spot"``, + ``"rect"``, ``"mesh"``). """ light: Light = env.sim.get_light(entity_cfg.uid) - num_instance = len(env_ids) + if light is None: + return + + is_global = light.is_global if hasattr(light, "is_global") else False + + # For global lights, normalize env_ids to avoid index-out-of-range + if is_global or light.num_instances == 1: + num_instance = 1 + effective_env_ids = None + else: + num_instance = len(env_ids) + effective_env_ids = env_ids if position_range: - init_pos = light.cfg.init_pos - new_pos = ( - torch.tensor(init_pos, dtype=torch.float32) - .unsqueeze_(0) - .repeat(num_instance, 1) - ) - random_value = sample_uniform( - lower=torch.tensor(position_range[0]), - upper=torch.tensor(position_range[1]), - size=new_pos.shape, - ) - new_pos += random_value - light.set_local_pose(new_pos, env_ids=env_ids) + if is_global: + logger.log_warning( + f"position_range ignored for global light '{entity_cfg.uid}' " + f"(type='{light.cfg.light_type}'). Use direction_range instead." + ) + else: + init_pos = light.cfg.init_pos + new_pos = ( + torch.tensor(init_pos, dtype=torch.float32) + .unsqueeze_(0) + .repeat(num_instance, 1) + ) + random_value = sample_uniform( + lower=torch.tensor(position_range[0]), + upper=torch.tensor(position_range[1]), + size=new_pos.shape, + ) + new_pos += random_value + light.set_local_pose(new_pos, env_ids=effective_env_ids) if color_range: color = torch.zeros((num_instance, 3), dtype=torch.float32) @@ -328,7 +356,7 @@ def randomize_light( size=color.shape, ) color += random_value - light.set_color(color, env_ids=env_ids) + light.set_color(color, env_ids=effective_env_ids) if intensity_range: init_intensity = light.cfg.intensity @@ -344,7 +372,31 @@ def randomize_light( ) new_intensity += random_value new_intensity.squeeze_(1) - light.set_intensity(new_intensity, env_ids=env_ids) + light.set_intensity(new_intensity, env_ids=effective_env_ids) + + if direction_range: + light_type = light.cfg.light_type + if light_type not in ("sun", "direction", "spot", "rect", "mesh"): + logger.log_warning( + f"direction_range ignored for light '{entity_cfg.uid}' " + f"(type='{light_type}'). Direction only applicable to " + f"'sun', 'direction', 'spot', 'rect', and 'mesh' types." + ) + return + + init_dir = light.cfg.direction + new_dir = ( + torch.tensor(init_dir, dtype=torch.float32) + .unsqueeze_(0) + .repeat(num_instance, 1) + ) + random_value = sample_uniform( + lower=torch.tensor(direction_range[0]), + upper=torch.tensor(direction_range[1]), + size=new_dir.shape, + ) + new_dir += random_value + light.set_direction(new_dir, env_ids=effective_env_ids) def randomize_emission_light( diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index e130bc6ef..84abb540a 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -18,6 +18,8 @@ import enum import json import os + +import dexsim import numpy as np import torch @@ -53,6 +55,122 @@ DEFAULT_RENDERER: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" +@configclass +class DLSSCfg: + """DLSS 3.5 configuration for the simulation window. + + NVIDIA DLSS 3.5 provides two AI-powered features for the OfflineRT renderer: + + - **Ray Reconstruction (RR)**: Replaces the OptiX denoiser. It takes the noisy + 1-spp path-traced HDR image plus G-buffer guides (albedo, normals, roughness, + depth, motion vectors) and reconstructs a clean image at render resolution. + - **Super Resolution (SR)**: Upscales the RR output from render resolution to + target/display resolution. + + When :attr:`dlss_enabled` is ``True``, both Ray Reconstruction and Super + Resolution are enabled together. The upscaling ratio is controlled through + :attr:`dlss_quality` (for example, ``5`` = DLAA / 1.0x, ``3`` = Quality). + + .. attention:: + DLSS is only available with the ``"rt"`` (OfflineRT) renderer in windowed + mode. Offscreen cameras always use the OptiX denoiser regardless of DLSS + settings. The window resolution (``SimulationManagerCfg.width/height``) + should match ``target_width/target_height`` when the target resolution is + explicitly set. + + Reference: ``developer_docs/dlss/README_DLSS.md`` in the dexsim repository. + """ + + dlss_enabled: bool = True + """Master DLSS enable toggle. Off → standard OptiX denoiser / TAA path.""" + + dlss_quality: int = -1 + """DLSS quality preset. ``-1`` = auto (ratio-based), ``0`` = UltraPerformance + (~3.0x), ``1`` = Performance (~2.25x), ``2`` = Balanced (~1.75x, default), + ``3`` = Quality (~1.5x), ``4`` = UltraQuality (~1.3x), ``5`` = DLAA (1.0x, + no upscale — render must equal target).""" + + upsample_ratio: float = 1.0 + """Convenience ratio ``target_width / render_width`` (and height). When + :attr:`render_width`/:attr:`render_height` are ``0``, the render resolution + is computed from the effective target resolution as + ``target / upsample_ratio``. Defaults to ``1.0`` (1:1 / DLAA mode). Must be + ``>= 1.0``. Explicit render or target dimensions take precedence over this + value.""" + + render_width: int = 0 + """Internal render resolution width. ``0`` = use window width. + The path tracer and DLSS-RR operate at this resolution.""" + + render_height: int = 0 + """Internal render resolution height. ``0`` = use window height.""" + + target_width: int = 0 + """Target/display resolution width. ``0`` = use window width. + DLSS-SR upscales to this resolution. Should match the window size.""" + + target_height: int = 0 + """Target/display resolution height. ``0`` = use window height.""" + + exposure_compensation: float = 1.0 + """Multiplier on the DLSS-RR pre-exposure scalar. ``>1`` (e.g. ``2.0``) + brightens the HDR fed to RR, speeding ghost rejection in dark areas; + ``<1`` darkens. ``1.0`` = neutral. Typical range 0.5–8.0; tune per scene. + No rebuild needed.""" + + def to_dexsim_cfg(self, window_width: int, window_height: int) -> dexsim.DLSSConfig: + """Resolve effective DLSS settings and return a dexsim config object. + + The effective target resolution defaults to the window size unless + :attr:`target_width`/:attr:`target_height` are set explicitly. The render + resolution defaults to the target unless :attr:`render_width`/ + :attr:`render_height` are explicit or :attr:`upsample_ratio` is provided. + + Args: + window_width: Window width in pixels. + window_height: Window height in pixels. + + Returns: + Populated :class:`dexsim.DLSSConfig` instance ready to assign to + ``world_config.dlss_config``. + """ + effective_target_width = ( + self.target_width if self.target_width > 0 else window_width + ) + effective_target_height = ( + self.target_height if self.target_height > 0 else window_height + ) + + if ( + self.upsample_ratio >= 1.0 + and self.render_width == 0 + and self.render_height == 0 + ): + effective_render_width = int(effective_target_width / self.upsample_ratio) + effective_render_height = int(effective_target_height / self.upsample_ratio) + else: + effective_render_width = ( + self.render_width if self.render_width > 0 else effective_target_width + ) + effective_render_height = ( + self.render_height + if self.render_height > 0 + else effective_target_height + ) + + dlss = dexsim.DLSSConfig() + dlss.dlss_enabled = self.dlss_enabled + dlss.rayreconstruction_enabled = True + dlss.upscale_enabled = True + dlss.dlss_quality = self.dlss_quality + dlss.render_width = effective_render_width + dlss.render_height = effective_render_height + dlss.target_width = effective_target_width + dlss.target_height = effective_target_height + dlss.exposure_compensation = self.exposure_compensation + return dlss + + @configclass class RenderCfg: renderer: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" @@ -71,6 +189,10 @@ class RenderCfg: spp: int = 1 """Samples per pixel for ray tracing rendering. This parameter is only valid when renderer is 'hybrid', 'fast-rt' or 'rt'.""" + dlss: DLSSCfg = field(default_factory=DLSSCfg) + """DLSS 3.5 configuration for AI-powered denoising and upscaling. + Only effective when ``renderer="rt"`` (OfflineRT) and in windowed mode.""" + def to_dexsim_flags(self): if self.renderer == "hybrid": return Renderer.HYBRID @@ -795,19 +917,108 @@ def from_dict(cls, init_dict: Dict[str, Union[str, float, tuple]]) -> ObjectBase class LightCfg(ObjectBaseCfg): """Configuration for a light asset in the simulation. - This class extends the base asset configuration to include specific properties for lights, + Supports six light types matching the dexsim rendering backend: + + - ``"point"``: Per-environment omnidirectional point light with position + and falloff radius. Created as a batched light (one per environment). + - ``"sun"``: Global directional sun light (infinite distance). Created as + a single scene-level instance. Uses direction only; position is ignored. + Sun-specific fields (``angular_radius``, ``halo_size``, ``halo_falloff``) + are reserved for future backend support. + - ``"direction"``: Global pure directional light at infinite distance. + Created as a single scene-level instance. Direction only; no position. + - ``"spot"``: Per-environment spotlight with position, direction, and + inner/outer cone angles. Created as a batched light. + - ``"rect"``: Per-environment rectangular area light with position, + direction, width, and height. Created as a batched light. + - ``"mesh"``: Per-environment mesh-based emissive light. Requires a + :class:`~dexsim.models.MeshObject` via + :meth:`embodichain.lab.sim.objects.light.Light.set_mesh` + (not tensor-batched). Created as a batched light. + + .. attention:: + The ``angular_radius``, ``halo_size``, and ``halo_falloff`` fields are + reserved for future use. The dexsim Python bindings do not yet expose + setters for these sun-specific properties. """ - # TODO: to be added more light type, such as spot, sun, etc. - light_type: Literal["point"] = "point" + light_type: Literal["point", "sun", "direction", "spot", "rect", "mesh"] = "point" + """Light type. Supported: ``"point"``, ``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, ``"mesh"``.""" + + # ------------------------------------------------------------------ + # Universal properties (apply to all light types) + # ------------------------------------------------------------------ color: tuple[float, float, float] = (1.0, 1.0, 1.0) + """RGB color of the light source. Defaults to white ``(1.0, 1.0, 1.0)``.""" + + intensity: float = 30.0 + """Intensity of the light source in watts/m^2. Defaults to ``30.0``.""" + + enable_shadow: bool = True + """Whether the light casts shadows. Defaults to ``True``.""" + + # ------------------------------------------------------------------ + # Point light + # ------------------------------------------------------------------ + + radius: float = 10.0 + """Falloff radius for point lights. Only used when ``light_type="point"``. Defaults to ``10.0``.""" + + # ------------------------------------------------------------------ + # Directional properties (sun, direction, spot, rect, mesh) + # ------------------------------------------------------------------ + + direction: tuple[float, float, float] = (0.0, 0.0, -1.0) + """Direction vector for directional, spot, rect, and mesh lights. + Defaults to ``(0.0, 0.0, -1.0)`` (pointing down along -Z).""" + + # ------------------------------------------------------------------ + # Sun light (reserved — Python bindings not yet available) + # ------------------------------------------------------------------ + + angular_radius: float = 0.5 + """Angular radius of the sun disc in degrees. Reserved for future use.""" + + halo_size: float = 10.0 + """Halo size for sun light. Reserved for future use.""" + + halo_falloff: float = 3.0 + """Halo falloff for sun light. Reserved for future use.""" + + # ------------------------------------------------------------------ + # Spot light + # ------------------------------------------------------------------ + + spot_angle_inner: float = 30.0 + """Inner cone angle of the spotlight in degrees. Only used when ``light_type="spot"``. + Defaults to ``30.0``.""" + + spot_angle_outer: float = 45.0 + """Outer cone angle of the spotlight in degrees. Only used when ``light_type="spot"``. + Defaults to ``45.0``.""" + + # ------------------------------------------------------------------ + # Rect light + # ------------------------------------------------------------------ + + rect_width: float = 1.0 + """Width of the rectangular area light. Only used when ``light_type="rect"``. + Defaults to ``1.0``.""" + + rect_height: float = 1.0 + """Height of the rectangular area light. Only used when ``light_type="rect"``. + Defaults to ``1.0``.""" - intensity: float = 50.0 - """Intensity of the light source with unit of watts/m^2.""" + # ------------------------------------------------------------------ + # Mesh light + # ------------------------------------------------------------------ - radius: float = 1e2 - """Falloff of the light, only used for point light.""" + mesh_path: str = "" + """Asset path for mesh-based emissive lights. Only used when ``light_type="mesh"``. + The actual mesh assignment is done via + :meth:`embodichain.lab.sim.objects.light.Light.set_mesh` which accepts a + :class:`dexsim.models.MeshObject`. This field stores the path for reference.""" @configclass diff --git a/embodichain/lab/sim/objects/light.py b/embodichain/lab/sim/objects/light.py index 8cea8d8ae..065267333 100644 --- a/embodichain/lab/sim/objects/light.py +++ b/embodichain/lab/sim/objects/light.py @@ -14,14 +14,19 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import numpy as np -from typing import List, Sequence +from typing import TYPE_CHECKING, List, Sequence from dexsim.render import Light as _Light from embodichain.lab.sim.cfg import LightCfg from embodichain.lab.sim.common import BatchEntity from embodichain.utils import logger +if TYPE_CHECKING: + from dexsim.models import MeshObject + class Light(BatchEntity): """Light represents a batch of lights in the simulation. @@ -87,6 +92,281 @@ def set_falloff( """ self._apply_scalar(falloffs, env_ids, "set_falloff") + def set_direction( + self, directions: torch.Tensor, env_ids: Sequence[int] | None = None + ) -> None: + """Set direction for directional-type lights. + + Only applies to ``sun``, ``direction``, ``spot``, ``rect``, and ``mesh`` + light types. Logs a warning and no-ops for other types. + + Args: + directions (torch.Tensor): Tensor of shape (3,) or (M, 3), representing + (x, y, z) direction vectors. + env_ids (Sequence[int] | None): Indices of instances to set. If None: + - For shape (3,), applies to all instances. + - For shape (M, 3), M must equal num_instances, applies per-instance. + """ + if self.cfg.light_type not in ("sun", "direction", "spot", "rect", "mesh"): + logger.log_warning( + f"set_direction not applicable to light type " + f"'{self.cfg.light_type}', ignoring." + ) + return + self._apply_vector3(directions, env_ids, "set_direction") + + def set_spot_angle( + self, + inner_angles: torch.Tensor, + outer_angles: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set inner and outer cone angles for spot lights. + + Only applies to ``spot`` light type. Logs a warning and no-ops for other types. + + Args: + inner_angles (torch.Tensor): Tensor of shape () (0-dim), (1,), or (M,) + representing inner cone angle in degrees. + outer_angles (torch.Tensor): Tensor of shape () (0-dim), (1,), or (M,) + representing outer cone angle in degrees. + env_ids (Sequence[int] | None): Indices of instances to set. + """ + if self.cfg.light_type != "spot": + logger.log_warning( + f"set_spot_angle not applicable to light type " + f"'{self.cfg.light_type}', ignoring." + ) + return + + if not torch.is_tensor(inner_angles) or not torch.is_tensor(outer_angles): + logger.log_error("set_spot_angle requires torch.Tensor arguments") + return + + inner_cpu = inner_angles.detach().cpu() + outer_cpu = outer_angles.detach().cpu() + + if env_ids is None: + all_ids = list(range(self.num_instances)) + else: + all_ids = list(env_ids) + + # Both scalar (0-dim): broadcast to all_ids + if inner_cpu.ndim == 0 and outer_cpu.ndim == 0: + iv = float(inner_cpu.item()) + ov = float(outer_cpu.item()) + for i in all_ids: + self._entities[i].set_spot_angle(iv, ov) + return + + # Both 1D with matching length + if inner_cpu.ndim == 1 and outer_cpu.ndim == 1: + ilen, olen = inner_cpu.shape[0], outer_cpu.shape[0] + inner_arr = inner_cpu.numpy() + outer_arr = outer_cpu.numpy() + + if ilen == olen == self.num_instances and env_ids is None: + for i in range(self.num_instances): + self._entities[i].set_spot_angle( + float(inner_arr[i]), float(outer_arr[i]) + ) + return + + if env_ids is not None and ilen == olen == len(all_ids): + for idx, i in enumerate(all_ids): + self._entities[i].set_spot_angle( + float(inner_arr[idx]), float(outer_arr[idx]) + ) + return + + # length-1 broadcast + if ilen == olen == 1: + iv = float(inner_arr[0]) + ov = float(outer_arr[0]) + for i in all_ids: + self._entities[i].set_spot_angle(iv, ov) + return + + logger.log_error( + f"set_spot_angle: invalid tensor shapes " + f"inner={tuple(inner_cpu.shape)}, outer={tuple(outer_cpu.shape)}" + ) + + def set_rect_wh( + self, + widths: torch.Tensor, + heights: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set width and height for rectangular area lights. + + Only applies to ``rect`` light type. Logs a warning and no-ops for other types. + + Args: + widths (torch.Tensor): Tensor of shape () (0-dim), (1,), or (M,) + representing width of the rectangular light. + heights (torch.Tensor): Tensor of shape () (0-dim), (1,), or (M,) + representing height of the rectangular light. + env_ids (Sequence[int] | None): Indices of instances to set. + """ + if self.cfg.light_type != "rect": + logger.log_warning( + f"set_rect_wh not applicable to light type " + f"'{self.cfg.light_type}', ignoring." + ) + return + + if not torch.is_tensor(widths) or not torch.is_tensor(heights): + logger.log_error("set_rect_wh requires torch.Tensor arguments") + return + + w_cpu = widths.detach().cpu() + h_cpu = heights.detach().cpu() + + if env_ids is None: + all_ids = list(range(self.num_instances)) + else: + all_ids = list(env_ids) + + # Both scalar (0-dim): broadcast to all_ids + if w_cpu.ndim == 0 and h_cpu.ndim == 0: + wv = float(w_cpu.item()) + hv = float(h_cpu.item()) + for i in all_ids: + self._entities[i].set_rect_wh(wv, hv) + return + + # Both 1D with matching length + if w_cpu.ndim == 1 and h_cpu.ndim == 1: + wlen, hlen = w_cpu.shape[0], h_cpu.shape[0] + w_arr = w_cpu.numpy() + h_arr = h_cpu.numpy() + + if wlen == hlen == self.num_instances and env_ids is None: + for i in range(self.num_instances): + self._entities[i].set_rect_wh(float(w_arr[i]), float(h_arr[i])) + return + + if env_ids is not None and wlen == hlen == len(all_ids): + for idx, i in enumerate(all_ids): + self._entities[i].set_rect_wh(float(w_arr[idx]), float(h_arr[idx])) + return + + # length-1 broadcast + if wlen == hlen == 1: + wv = float(w_arr[0]) + hv = float(h_arr[0]) + for i in all_ids: + self._entities[i].set_rect_wh(wv, hv) + return + + logger.log_error( + f"set_rect_wh: invalid tensor shapes " + f"width={tuple(w_cpu.shape)}, height={tuple(h_cpu.shape)}" + ) + + def set_mesh( + self, + mesh: "MeshObject", + env_ids: Sequence[int] | None = None, + ) -> None: + """Set the mesh for mesh-type lights. + + Only applies to ``mesh`` light type. Logs a warning and no-ops for other types. + This is NOT tensor-batched — the same MeshObject is assigned to all targeted + instances. + + Args: + mesh (MeshObject): The mesh object to assign to the light. + env_ids (Sequence[int] | None): Indices of instances to set. If None, + applies to all instances. + """ + if self.cfg.light_type != "mesh": + logger.log_warning( + f"set_mesh not applicable to light type " + f"'{self.cfg.light_type}', ignoring." + ) + return + + if env_ids is None: + target_ids = list(range(self.num_instances)) + else: + target_ids = list(env_ids) + + for i in target_ids: + try: + self._entities[i].set_mesh(mesh) + except Exception as e: + logger.log_error(f"set_mesh: error for instance {i}: {e}") + + def enable_shadow( + self, + flags: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Enable or disable shadow casting. + + Applies to all light types. + + Args: + flags (torch.Tensor): Boolean tensor of shape () (0-dim), (1,), or (M,). + Non-zero values enable shadows; zero disables. + env_ids (Sequence[int] | None): Indices of instances to set. + """ + if not torch.is_tensor(flags): + logger.log_error( + f"enable_shadow requires a torch.Tensor, got {type(flags)}" + ) + return + + cpu = flags.detach().cpu() + if env_ids is None: + all_ids = list(range(self.num_instances)) + else: + all_ids = list(env_ids) + + # Scalar: broadcast + if cpu.ndim == 0: + val = bool(cpu.item() != 0) + for i in all_ids: + self._entities[i].set_shadow(val) + return + + # 1D tensor + if cpu.ndim == 1: + length = cpu.shape[0] + arr = cpu.numpy() + if length == self.num_instances and env_ids is None: + for i in range(self.num_instances): + self._entities[i].set_shadow(bool(arr[i] != 0)) + return + if env_ids is not None and length == len(all_ids): + for idx, i in enumerate(all_ids): + self._entities[i].set_shadow(bool(arr[idx] != 0)) + return + if length == 1: + val = bool(arr[0] != 0) + for i in all_ids: + self._entities[i].set_shadow(val) + return + + logger.log_error( + f"enable_shadow: tensor shape {tuple(cpu.shape)} is invalid for broadcasting" + ) + + @property + def is_global(self) -> bool: + """Whether this light is a global scene light (single instance). + + Global lights (``"sun"``, ``"direction"``) are infinite-distance + light sources that illuminate the entire scene. They are never + batched per environment. + + Returns: + bool: True if the light is a global scene light. + """ + return self.cfg.light_type in ("sun", "direction") + def set_local_pose( self, pose: torch.Tensor, @@ -104,6 +384,13 @@ def set_local_pose( - For matrix input (4,4) broadcast to all, or (M,4,4) with M == num_instances. to_matrix (bool): Interpret `pose` as full 4x4 matrix if True, else as vector(s). """ + if self.is_global: + logger.log_warning( + f"set_local_pose not applicable to '{self.cfg.light_type}' light type " + f"(infinite distance, direction only). Use set_direction() instead." + ) + return + if not torch.is_tensor(pose): logger.log_error( f"set_local_pose requires a torch.Tensor, got {type(pose)}" @@ -292,8 +579,63 @@ def _apply_scalar( ) def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset the light to its initial configuration state. + + Applies only the properties relevant to ``self.cfg.light_type``. + + .. attention:: + When this light has only one instance (global lights such as + ``"sun"`` and ``"direction"``, or any light in a single-environment + scene), ``env_ids`` is normalized to ``None`` because there is only + one instance to update. Passing per-environment indices is silently + ignored. + + Args: + env_ids (Sequence[int] | None): The environment IDs to reset. + If None, resets all environments. + """ self.cfg: LightCfg + light_type = self.cfg.light_type + + # Global lights have a single instance — ignore per-env indices. + if self.num_instances == 1: + env_ids = None + + # Universal properties self.set_color(torch.as_tensor(self.cfg.color), env_ids=env_ids) self.set_intensity(torch.as_tensor(self.cfg.intensity), env_ids=env_ids) - self.set_falloff(torch.as_tensor(self.cfg.radius), env_ids=env_ids) - self.set_local_pose(torch.as_tensor(self.cfg.init_pos), env_ids=env_ids) + self.enable_shadow( + torch.as_tensor(float(self.cfg.enable_shadow)), env_ids=env_ids + ) + + # Position (all types except global infinite-distance lights) + if not self.is_global: + self.set_local_pose(torch.as_tensor(self.cfg.init_pos), env_ids=env_ids) + + # Point light: falloff + if light_type == "point": + self.set_falloff(torch.as_tensor(self.cfg.radius), env_ids=env_ids) + + # Directional types: direction vector + if light_type in ("sun", "direction", "spot", "rect", "mesh"): + self.set_direction(torch.as_tensor(self.cfg.direction), env_ids=env_ids) + + # Spot light: cone angles + if light_type == "spot": + self.set_spot_angle( + torch.as_tensor(self.cfg.spot_angle_inner), + torch.as_tensor(self.cfg.spot_angle_outer), + env_ids=env_ids, + ) + + # Rect light: dimensions + if light_type == "rect": + self.set_rect_wh( + torch.as_tensor(self.cfg.rect_width), + torch.as_tensor(self.cfg.rect_height), + env_ids=env_ids, + ) + + # Mesh light: mesh_path is stored in cfg but actual mesh assignment + # is done via set_mesh() which requires a MeshObject. + # Sun-specific angular_radius/halo are reserved for future backend support. diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index c14d98c99..cef45899f 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -75,6 +75,7 @@ ContactSensor, ) from embodichain.lab.sim.cfg import ( + DLSSCfg, RenderCfg, PhysicsCfg, MarkerCfg, @@ -96,6 +97,7 @@ __all__ = [ "SimulationManager", "SimulationManagerCfg", + "DLSSCfg", "SIM_CACHE_DIR", "MATERIAL_CACHE_DIR", "CONVEX_DECOMP_DIR", @@ -306,6 +308,7 @@ def __init__( self._create_default_plane() self.set_default_background() + self.set_default_global_lighting() # Set physics to manual update mode by default. self.set_manual_update(True) @@ -473,6 +476,37 @@ def _convert_sim_config( sim_config.render_cfg.spp ) + # Configure DLSS 3.5 (Ray Reconstruction + Super Resolution). + # Only effective with OfflineRT renderer in windowed mode. + dlss_cfg = sim_config.render_cfg.dlss + if dlss_cfg.dlss_enabled: + if sim_config.render_cfg.renderer != "rt": + logger.log_warning( + f"DLSS is enabled but renderer is '{sim_config.render_cfg.renderer}', " + f"not 'rt' (OfflineRT). DLSS only works with the OfflineRT renderer. " + f"DLSS settings will be ignored." + ) + else: + world_config.dlss_config = dlss_cfg.to_dexsim_cfg( + window_width=sim_config.width, + window_height=sim_config.height, + ) + dlss = world_config.dlss_config + world_config.raytrace_config.window_taa_enabled = False + + # Align the window size with the effective target resolution. + # Rule: win_config.width/height == target_width/target_height. + win_config.width = dlss.target_width + win_config.height = dlss.target_height + + logger.log_info( + f"DLSS enabled with renderer='{sim_config.render_cfg.renderer}': " + f"render={dlss.render_width}x{dlss.render_height}, " + f"target={dlss.target_width}x{dlss.target_height}, " + f"upsample_ratio={dlss_cfg.upsample_ratio}, " + f"quality={dlss.dlss_quality}." + ) + if type(sim_config.sim_device) is str: self.device = torch.device(sim_config.sim_device) else: @@ -701,6 +735,28 @@ def _create_default_plane(self): # TODO: add default physics attributes for the plane. + def set_default_global_lighting(self) -> None: + """Set default global lighting for the scene. + + Configures both the environment emission (ambient) light and a + directional light to provide default scene illumination. The + directional light is a global scene light (infinite distance) + pointing downward along the -Z axis. + """ + # Environment emission light + self.set_emission_light([1.0, 1.0, 1.0], 120.0) + + # Directional light as global scene light + dir_light_cfg = LightCfg( + uid="default_global_light", + light_type="sun", + intensity=8.0, + direction=(0.0, 0.0, -1.0), + color=(1.0, 0.95, 0.85), + enable_shadow=True, + ) + self.add_light(dir_light_cfg) + def set_default_background(self) -> None: """Set default background.""" @@ -714,12 +770,10 @@ def set_default_background(self) -> None: uid=mat_name, base_color_texture=color_texture, roughness_texture=roughness_texture, - roughness=0.7, + roughness=1.0, ) ) - self.set_emission_light([1.0, 1.0, 1.0], 120.0) - self._default_plane.set_material(mat.get_instance("plane_mat").mat) self._visual_materials[mat_name] = mat @@ -797,14 +851,42 @@ def get_asset( logger.log_warning(f"Asset {uid} not found.") return None + # Light type string → dexsim LightType enum mapping + _LIGHT_TYPE_MAP: dict[str, LightType] = { + "point": LightType.POINT, + "sun": LightType.SUN, + "direction": LightType.DIRECTION, + "spot": LightType.SPOT, + "rect": LightType.RECT, + "mesh": LightType.MESH, + } + + # Light types that are created as a single global scene light (not per-environment). + _GLOBAL_LIGHT_TYPES: tuple[str, ...] = ("sun", "direction") + def add_light(self, cfg: LightCfg) -> Light: """Create a light in the scene. + Supports six light types: ``"point"``, ``"sun"``, ``"direction"``, + ``"spot"``, ``"rect"``, and ``"mesh"``. See :class:`LightCfg` for + type-specific configuration fields. + + .. attention:: + ``"sun"`` and ``"direction"`` lights are global scene lights + (infinite-distance directional light sources). They are created + as a single instance on the root environment, not batched per + environment. All other types are created as per-environment + batched lights. + Args: - cfg (LightCfg): Configuration for the light, including type, color, intensity, and radius. + cfg (LightCfg): Configuration for the light, including type, color, + intensity, and type-specific properties. Returns: Light: The created light instance. + + Raises: + RuntimeError: If ``cfg.light_type`` is not one of the supported types. """ if cfg.uid is None: uid = "light" @@ -815,22 +897,42 @@ def add_light(self, cfg: LightCfg) -> Light: if uid in self._lights: logger.log_error(f"Light {uid} already exists.") - light_type = cfg.light_type - if light_type == "point": - light_type = LightType.POINT - else: + light_type_str = cfg.light_type + light_type = self._LIGHT_TYPE_MAP.get(light_type_str) + if light_type is None: + supported = ", ".join(self._LIGHT_TYPE_MAP.keys()) logger.log_error( - f"Unsupported light type: {light_type}. Supported types: point." + f"Unsupported light type: '{light_type_str}'. " + f"Supported types: {supported}." ) - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - light_list = [] - for i, env in enumerate(env_list): - light_name = f"{uid}_{i}" - light = env.create_light(light_name, light_type) - light_list.append(light) + # Validation warnings for type-specific constraints + if light_type_str == "mesh" and not cfg.mesh_path: + logger.log_warning( + f"Mesh light '{uid}' has no mesh_path set. " + f"Use set_mesh() to assign a MeshObject." + ) + if light_type_str == "rect" and (cfg.rect_width <= 0 or cfg.rect_height <= 0): + logger.log_warning( + f"Rect light '{uid}' has zero or negative dimensions " + f"(width={cfg.rect_width}, height={cfg.rect_height})." + ) - batch_lights = Light(cfg=cfg, entities=light_list) + if cfg.light_type in self._GLOBAL_LIGHT_TYPES: + # Global scene light: create a single instance on the root + # environment. Infinite-distance lights (sun, direction) are + # physically scene-global and should not be duplicated per arena. + light = self._env.create_light(uid, light_type) + batch_lights = Light(cfg=cfg, entities=[light]) + else: + # Per-environment batched light: one instance per arena. + env_list = [self._env] if len(self._arenas) == 0 else self._arenas + light_list = [] + for i, env in enumerate(env_list): + light_name = f"{uid}_{i}" + light = env.create_light(light_name, light_type) + light_list.append(light) + batch_lights = Light(cfg=cfg, entities=light_list) self._lights[uid] = batch_lights diff --git a/tests/gym/envs/managers/test_event_functors.py b/tests/gym/envs/managers/test_event_functors.py index 981e44ab0..4c6cb8f06 100644 --- a/tests/gym/envs/managers/test_event_functors.py +++ b/tests/gym/envs/managers/test_event_functors.py @@ -117,6 +117,64 @@ def set_mass(self, link_name: str, mass: float) -> int: return 0 +class MockLight: + """Mock light for event functor tests. + + Supports both global lights (num_instances=1 for sun/direction) + and per-environment batched lights (num_instances=num_envs). + """ + + def __init__( + self, + uid: str = "test_light", + num_envs: int = 4, + light_type: str = "point", + init_pos: tuple = (0.0, 0.0, 2.0), + direction: tuple = (0.0, 0.0, -1.0), + intensity: float = 30.0, + color: tuple = (1.0, 1.0, 1.0), + ): + self.uid = uid + + # Config-like attributes + class CfgProxy: + pass + + self.cfg = CfgProxy() + self.cfg.light_type = light_type + self.cfg.init_pos = init_pos + self.cfg.direction = direction + self.cfg.intensity = intensity + self.cfg.color = color + + self.device = torch.device("cpu") + self.is_global = light_type in ("sun", "direction") + self.num_instances = 1 if self.is_global else num_envs + + # Track calls for assertions + self._last_set_color = None + self._last_set_intensity = None + self._last_set_local_pose = None + self._last_set_direction = None + self._last_env_ids = None + + def set_color(self, color, env_ids=None): + self._last_set_color = color + self._last_env_ids = env_ids + + def set_intensity(self, intensity, env_ids=None): + self._last_set_intensity = intensity + self._last_env_ids = env_ids + + def set_local_pose(self, pose, env_ids=None, to_matrix=False): + self._last_set_local_pose = pose + self._last_env_ids = env_ids + + def set_direction(self, direction, env_ids=None): + self._last_set_direction = direction + self._last_env_ids = env_ids + + class MockArticulation: """Mock articulation for event functor tests.""" @@ -238,6 +296,7 @@ def __init__(self, num_envs: int = 4): self._articulations = {} self._robots = {} self._rigid_object_groups = {} + self._lights = {} def get_rigid_object(self, uid: str): return self._rigid_objects.get(uid) @@ -290,6 +349,13 @@ def set_emission_light(self, color=None, intensity=None) -> None: self._last_emission_color = color self._last_emission_intensity = intensity + def get_light(self, uid: str): + return self._lights.get(uid) + + def add_light(self, light: MockLight): + self._lights[light.uid] = light + return light + class MockEnv: """Mock environment for event functor tests.""" @@ -333,7 +399,9 @@ def __init__(self, num_envs: int = 4, num_joints: int = 6): ) from embodichain.lab.gym.envs.managers.randomization.visual import ( randomize_indirect_lighting, + randomize_light, ) +from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg from embodichain.lab.gym.envs.managers import FunctorCfg @@ -1010,3 +1078,260 @@ def test_emissive_values_vary_across_calls(self): # Over 20 draws from [0, 1000] all values being identical is astronomically unlikely assert len(intensities) > 1 + + +# -------------------------------------------------------------------------- +# randomize_light tests +# -------------------------------------------------------------------------- + + +class TestRandomizeLight: + """Tests for the randomize_light function.""" + + # ------------------------------------------------------------------ + # Per-environment batched light tests (point light) + # ------------------------------------------------------------------ + + def test_randomize_color(self): + """color_range randomizes the light color for all envs.""" + env = MockEnv(num_envs=4) + light = MockLight("test_pt", num_envs=4, light_type="point") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_pt") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + color_range=[[0.2, 0.2, 0.2], [0.8, 0.8, 0.8]], + ) + + assert light._last_set_color is not None + assert light._last_set_color.shape == (4, 3) + # All values should be within [0.2, 0.8] + assert (light._last_set_color >= 0.2).all() + assert (light._last_set_color <= 0.8).all() + + def test_randomize_intensity(self): + """intensity_range randomizes the light intensity for all envs. + + The intensity_range is additive: random value is added to cfg.intensity. + """ + env = MockEnv(num_envs=4) + light = MockLight("test_pt", num_envs=4, light_type="point") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_pt") + env_ids = torch.tensor([0, 1, 2, 3]) + + # Use a small additive range so we can bound the result + randomize_light( + env, + env_ids, + entity_cfg, + intensity_range=(10.0, 20.0), + ) + + assert light._last_set_intensity is not None + assert light._last_set_intensity.shape == (4,) + # init_intensity=30 + rand(10,20) => [40, 50] + assert (light._last_set_intensity >= 40.0).all() + assert (light._last_set_intensity <= 50.0).all() + + def test_randomize_position(self): + """position_range randomizes the light position for batched lights.""" + env = MockEnv(num_envs=4) + light = MockLight("test_pt", num_envs=4, light_type="point") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_pt") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + position_range=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], + ) + + assert light._last_set_local_pose is not None + assert light._last_set_local_pose.shape == (4, 3) + + def test_randomize_direction(self): + """direction_range randomizes the light direction for spot lights.""" + env = MockEnv(num_envs=4) + light = MockLight("test_spot", num_envs=4, light_type="spot") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_spot") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + direction_range=[[-0.2, -0.2, -0.2], [0.2, 0.2, 0.2]], + ) + + assert light._last_set_direction is not None + assert light._last_set_direction.shape == (4, 3) + + def test_direction_range_ignored_for_point_light(self): + """direction_range is ignored for point lights (no crash).""" + env = MockEnv(num_envs=4) + light = MockLight("test_pt", num_envs=4, light_type="point") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_pt") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + direction_range=[[-0.1, -0.1, -0.1], [0.1, 0.1, 0.1]], + ) + + # direction should NOT have been set for point light + assert light._last_set_direction is None + + def test_combined_randomization(self): + """All four ranges can be combined in one call.""" + env = MockEnv(num_envs=4) + light = MockLight("test_spot", num_envs=4, light_type="spot") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_spot") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + position_range=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], + color_range=[[0.2, 0.2, 0.2], [0.8, 0.8, 0.8]], + intensity_range=(50.0, 100.0), + direction_range=[[-0.1, -0.1, -0.1], [0.1, 0.1, 0.1]], + ) + + assert light._last_set_local_pose is not None + assert light._last_set_color is not None + assert light._last_set_intensity is not None + assert light._last_set_direction is not None + + def test_light_not_found_returns_early(self): + """No crash when light UID is not registered.""" + env = MockEnv(num_envs=4) + entity_cfg = SceneEntityCfg(uid="nonexistent") + env_ids = torch.tensor([0, 1, 2, 3]) + + # Should not raise + try: + randomize_light( + env, + env_ids, + entity_cfg, + color_range=[[0.2, 0.2, 0.2], [0.8, 0.8, 0.8]], + ) + except Exception as e: + pytest.fail(f"randomize_light should not crash on missing light: {e}") + + # ------------------------------------------------------------------ + # Global scene light tests (sun, direction) + # ------------------------------------------------------------------ + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_global_light_position_range_ignored(self, light_type): + """position_range is ignored for global lights (sun/direction).""" + env = MockEnv(num_envs=4) + light = MockLight("test_global", num_envs=1, light_type=light_type) + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_global") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + position_range=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], + ) + + # set_local_pose should NOT be called for global lights + assert light._last_set_local_pose is None + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_global_light_color_applied_single_instance(self, light_type): + """color_range works for global lights — single instance broadcast.""" + env = MockEnv(num_envs=4) + light = MockLight("test_global", num_envs=1, light_type=light_type) + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_global") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + color_range=[[0.3, 0.3, 0.3], [0.7, 0.7, 0.7]], + ) + + assert light._last_set_color is not None + # Global light: single instance, color shape should be (1, 3) + assert light._last_set_color.shape == (1, 3) + # env_ids=None for single-instance lights + assert light._last_env_ids is None + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_global_light_intensity_applied_single_instance(self, light_type): + """intensity_range works for global lights — single instance.""" + env = MockEnv(num_envs=4) + light = MockLight("test_global", num_envs=1, light_type=light_type) + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_global") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + intensity_range=(50.0, 100.0), + ) + + assert light._last_set_intensity is not None + assert light._last_set_intensity.shape == (1,) + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_global_light_direction_range_applied(self, light_type): + """direction_range works for global lights (sun/direction).""" + env = MockEnv(num_envs=4) + light = MockLight("test_global", num_envs=1, light_type=light_type) + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_global") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light( + env, + env_ids, + entity_cfg, + direction_range=[[-0.2, -0.2, -0.2], [0.2, 0.2, 0.2]], + ) + + assert light._last_set_direction is not None + assert light._last_set_direction.shape == (1, 3) + assert light._last_env_ids is None + + # ------------------------------------------------------------------ + # Edge cases + # ------------------------------------------------------------------ + + def test_no_ranges_does_nothing(self): + """Calling with no ranges set does not crash.""" + env = MockEnv(num_envs=4) + light = MockLight("test_pt", num_envs=4, light_type="point") + env.sim.add_light(light) + entity_cfg = SceneEntityCfg(uid="test_pt") + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_light(env, env_ids, entity_cfg) + + # Nothing should have been called + assert light._last_set_color is None + assert light._last_set_intensity is None + assert light._last_set_local_pose is None + assert light._last_set_direction is None diff --git a/tests/sim/objects/test_light.py b/tests/sim/objects/test_light.py index 7e9d58c49..0b7bbd794 100644 --- a/tests/sim/objects/test_light.py +++ b/tests/sim/objects/test_light.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import pytest import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg @@ -31,7 +33,7 @@ def setup_method(self): "light_type": "point", "color": [0.1, 0.1, 0.1], "radius": 10.0, - "position": [0.0, 0.0, 2.0], + "init_pos": [0.0, 0.0, 2.0], "uid": "point_light", } self.light = self.sim.add_light(cfg=LightCfg.from_dict(cfg_dict)) @@ -159,3 +161,297 @@ def teardown_method(self): import gc gc.collect() + + +class TestLightTypes: + """Tests for all six supported light types.""" + + @pytest.fixture(autouse=True) + def setup(self): + """Create a SimulationManager for each test.""" + config = SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=4) + self.sim = SimulationManager(config) + yield + self.sim.destroy() + import embodichain.lab.sim as om + + om.SimulationManager.flush_cleanup_queue() + self.__dict__.clear() + import gc + + gc.collect() + + # ------------------------------------------------------------------ + # Creation tests + # ------------------------------------------------------------------ + + @pytest.mark.parametrize( + "light_type, expected_num_instances", + [ + ("point", 4), + ("sun", 1), + ("direction", 1), + ("spot", 4), + ("rect", 4), + ("mesh", 4), + ], + ) + def test_create_each_light_type(self, light_type, expected_num_instances): + """Each of the 6 light types can be created without error. + + ``sun`` and ``direction`` are global scene lights with a single instance. + All other types are per-environment batched lights. + """ + cfg = LightCfg( + uid=f"test_{light_type}", + light_type=light_type, + init_pos=(0.0, 0.0, 3.0), + ) + light = self.sim.add_light(cfg=cfg) + assert light is not None, f"Failed to create light of type '{light_type}'" + assert light.num_instances == expected_num_instances + if light_type in ("sun", "direction"): + assert light.is_global, f"{light_type} should be a global light" + + def test_unknown_light_type_errors(self): + """Passing an invalid light_type raises RuntimeError.""" + cfg = LightCfg(uid="bad", light_type="invalid") + with pytest.raises(RuntimeError, match="Unsupported light type"): + self.sim.add_light(cfg=cfg) + + def test_mesh_light_empty_path_warns(self): + """Mesh light with empty mesh_path warns at creation but succeeds.""" + cfg = LightCfg( + uid="mesh_no_path", + light_type="mesh", + init_pos=(0.0, 0.0, 2.0), + ) + light = self.sim.add_light(cfg=cfg) + assert light is not None, "Mesh light should be created even without path" + + # ------------------------------------------------------------------ + # Setter type-validation tests + # ------------------------------------------------------------------ + + def test_set_direction_on_point_warns(self): + """Calling set_direction on a point light logs a warning and no-ops.""" + cfg = LightCfg(uid="pt", light_type="point", init_pos=(0, 0, 2)) + light = self.sim.add_light(cfg=cfg) + direction = torch.tensor([1.0, 0.0, 0.0]) + # Should not raise; should log a warning + try: + light.set_direction(direction) + except Exception as e: + pytest.fail(f"set_direction on point light should warn, not crash: {e}") + + def test_set_spot_angle_on_point_warns(self): + """Calling set_spot_angle on a non-spot light warns and no-ops.""" + cfg = LightCfg(uid="pt2", light_type="point", init_pos=(0, 0, 2)) + light = self.sim.add_light(cfg=cfg) + inner = torch.tensor(15.0) + outer = torch.tensor(30.0) + try: + light.set_spot_angle(inner, outer) + except Exception as e: + pytest.fail(f"set_spot_angle on point light should warn, not crash: {e}") + + def test_set_rect_wh_on_point_warns(self): + """Calling set_rect_wh on a non-rect light warns and no-ops.""" + cfg = LightCfg(uid="pt3", light_type="point", init_pos=(0, 0, 2)) + light = self.sim.add_light(cfg=cfg) + w = torch.tensor(2.0) + h = torch.tensor(3.0) + try: + light.set_rect_wh(w, h) + except Exception as e: + pytest.fail(f"set_rect_wh on point light should warn, not crash: {e}") + + def test_set_local_pose_on_direction_warns(self): + """Calling set_local_pose on a direction light warns and no-ops.""" + cfg = LightCfg( + uid="dir_light", + light_type="direction", + direction=(0.0, -1.0, 0.0), + ) + light = self.sim.add_light(cfg=cfg) + pose = torch.tensor([1.0, 2.0, 3.0]) + try: + light.set_local_pose(pose) + except Exception as e: + pytest.fail( + f"set_local_pose on direction light should warn, not crash: {e}" + ) + + def test_set_local_pose_on_sun_warns(self): + """Calling set_local_pose on a sun light warns and no-ops.""" + cfg = LightCfg( + uid="sun_light", + light_type="sun", + direction=(0.0, -1.0, 0.0), + ) + light = self.sim.add_light(cfg=cfg) + pose = torch.tensor([1.0, 2.0, 3.0]) + try: + light.set_local_pose(pose) + except Exception as e: + pytest.fail(f"set_local_pose on sun light should warn, not crash: {e}") + + # ------------------------------------------------------------------ + # Type-specific property tests + # ------------------------------------------------------------------ + + def test_spot_light_has_cone_angles(self): + """Spot light creation applies inner/outer cone angles from cfg.""" + cfg = LightCfg( + uid="spot_test", + light_type="spot", + init_pos=(0.0, 0.0, 3.0), + direction=(0.0, 0.0, -1.0), + spot_angle_inner=20.0, + spot_angle_outer=40.0, + ) + light = self.sim.add_light(cfg=cfg) + assert light is not None + # Should not crash on creation; spot_angle applied during reset() + + def test_rect_light_has_dimensions(self): + """Rect light creation applies width/height from cfg.""" + cfg = LightCfg( + uid="rect_test", + light_type="rect", + init_pos=(0.0, 0.0, 3.0), + direction=(0.0, 0.0, -1.0), + rect_width=2.0, + rect_height=1.5, + ) + light = self.sim.add_light(cfg=cfg) + assert light is not None + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_global_light_no_position_in_reset(self, light_type): + """Global light (sun/direction) reset does not set position.""" + cfg = LightCfg( + uid=f"{light_type}_test", + light_type=light_type, + direction=(0.0, -1.0, 0.0), + ) + light = self.sim.add_light(cfg=cfg) + assert light is not None + assert light.is_global + assert light.num_instances == 1 + + @pytest.mark.parametrize("light_type", ["sun", "direction"]) + def test_reset_global_light_with_env_ids(self, light_type): + """reset_objects_state with env_ids does not crash for global lights.""" + cfg = LightCfg( + uid=f"global_{light_type}", + light_type=light_type, + direction=(0.0, -1.0, 0.0), + ) + self.sim.add_light(cfg=cfg) + try: + self.sim.reset_objects_state(env_ids=[1, 2]) + except Exception as e: + pytest.fail( + f"reset_objects_state with global light ({light_type}) failed: {e}" + ) + + # ------------------------------------------------------------------ + # Broadcasting tests + # ------------------------------------------------------------------ + + def test_set_direction_broadcasting(self): + """set_direction with (3,) tensor broadcasts to all instances.""" + cfg = LightCfg( + uid="spot_broadcast", + light_type="spot", + init_pos=(0, 0, 3), + direction=(0.0, 0.0, -1.0), + ) + light = self.sim.add_light(cfg=cfg) + new_dir = torch.tensor([1.0, 0.0, 0.0]) + try: + light.set_direction(new_dir) + except Exception as e: + pytest.fail(f"set_direction broadcast failed: {e}") + + def test_set_direction_with_env_ids(self): + """set_direction with (M, 3) tensor applies per-instance.""" + cfg = LightCfg( + uid="spot_env", + light_type="spot", + init_pos=(0, 0, 3), + direction=(0.0, 0.0, -1.0), + ) + light = self.sim.add_light(cfg=cfg) + env_ids = [0, 2] + dirs = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + try: + light.set_direction(dirs, env_ids=env_ids) + except Exception as e: + pytest.fail(f"set_direction per-instance failed: {e}") + + def test_enable_shadow(self): + """enable_shadow sets shadow flag on all instances.""" + cfg = LightCfg(uid="shadow_test", light_type="point", init_pos=(0, 0, 2)) + light = self.sim.add_light(cfg=cfg) + try: + light.enable_shadow(torch.tensor(0.0)) # disable + light.enable_shadow(torch.tensor(1.0)) # enable + except Exception as e: + pytest.fail(f"enable_shadow failed: {e}") + + # ------------------------------------------------------------------ + # from_dict compatibility + # ------------------------------------------------------------------ + + def test_from_dict_new_types(self): + """LightCfg.from_dict works with new type fields.""" + cfg = LightCfg.from_dict( + { + "light_type": "spot", + "color": [1.0, 0.9, 0.8], + "intensity": 100.0, + "init_pos": [0.0, 0.0, 3.0], + "direction": [0.0, 0.0, -1.0], + "spot_angle_inner": 25.0, + "spot_angle_outer": 50.0, + "uid": "from_dict_spot", + } + ) + assert cfg.light_type == "spot" + assert cfg.spot_angle_inner == 25.0 + assert cfg.spot_angle_outer == 50.0 + assert cfg.direction == (0.0, 0.0, -1.0) or cfg.direction == [0.0, 0.0, -1.0] + + # ------------------------------------------------------------------ + # Backward compatibility + # ------------------------------------------------------------------ + + def test_point_light_backward_compat(self): + """Existing point-light config pattern still works.""" + cfg_dict = { + "light_type": "point", + "color": [0.1, 0.1, 0.1], + "radius": 10.0, + "init_pos": [0.0, 0.0, 2.0], + "uid": "point_compat", + } + cfg = LightCfg.from_dict(cfg_dict) + assert cfg.init_pos == [0.0, 0.0, 2.0] + light = self.sim.add_light(cfg=cfg) + assert light is not None + assert light.num_instances == 4 + + # set_color should still work + base_color = torch.tensor([0.1, 0.1, 0.1]) + try: + light.set_color(base_color) + except Exception as e: + pytest.fail(f"set_color failed: {e}") + + # set_falloff should still work + try: + light.set_falloff(torch.tensor(100.0)) + except Exception as e: + pytest.fail(f"set_falloff failed: {e}")