# ----------------------------------------------------------------------------
# Copyright (c) 2021-2025 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 copy
import torch
import dexsim
import numpy as np
from typing import Optional, Dict, Union
from functools import cached_property
from dexsim.engine import MaterialInst, Material
from embodichain.lab.sim.utility import is_rt_enabled
from embodichain.utils import configclass, logger
[docs]
@configclass
class VisualMaterialCfg:
"""Configuration for visual material with PBR properties for rasterization and ray tracing."""
uid: str = "default_mat"
# Basic PBR properties
base_color: list = [0.5, 0.5, 0.5, 1.0]
"""Base color/diffuse color (RGBA)"""
metallic: float = 0.0
"""Metallic factor (0.0 = dielectric, 1.0 = metallic)"""
roughness: float = 0.5
"""Surface roughness (0.0 = smooth, 1.0 = rough)"""
# Additional PBR properties
emissive: list = [0.0, 0.0, 0.0] # Emissive color (RGB)
emissive_intensity: float = 1.0 # Emissive intensity multiplier
# Texture maps
base_color_texture: str = None
"""Base color texture map"""
metallic_texture: str = None
"""Metallic map"""
roughness_texture: str = None
"""Roughness map"""
normal_texture: str = None
"""Normal map"""
ao_texture: str = None
"""Ambient occlusion map"""
# Ray tracing specific properties
ior: float = 1.5
"""Index of refraction for ray tracing materials"""
rt_material_type: str = "BRDF_GGX_SMITH"
"""Ray tracing material type. Options: 'BRDF_GGX_SMITH', 'BTDF_GGX_SMITH', 'BSDF_GGX_SMITH'"""
# Currently disabled properties
# subsurface: float = 0.0 # Subsurface scattering factor
# subsurface_color: list = [1.0, 1.0, 1.0] # Subsurface scattering color
[docs]
@classmethod
def from_dict(cls, cfg_dict: dict) -> VisualMaterialCfg:
base = cls()
for k, v in cfg_dict.items():
if hasattr(base, k):
setattr(base, k, v)
else:
logger.log_warning(f"Unknown field '{k}' in VisualMaterialCfg.")
return base
[docs]
class VisualMaterial:
"""Visual material definition in the simulation environment.
A visual material is actually a material template from which material instances can be created.
It holds multiple material instances, which is used to assign to different objects in the environment.
"""
RT_MATERIAL_TYPES = [
"BRDF_GGX_SMITH",
"BTDF_GGX_SMITH",
"BSDF_GGX_SMITH",
]
[docs]
def __init__(self, cfg: VisualMaterialCfg, mat: Material):
self.uid = cfg.uid
self.cfg = copy.deepcopy(cfg)
self._mat = mat
self._default_mat_inst = self.create_instance(self.uid)
@cached_property
def is_rt_enabled(self) -> bool:
return is_rt_enabled()
@property
def mat(self) -> Material:
return self._mat
[docs]
def set_default_properties(
self, mat_inst: VisualMaterialInst, cfg: VisualMaterialCfg
) -> None:
mat_inst.set_base_color(cfg.base_color)
mat_inst.set_metallic(cfg.metallic)
mat_inst.set_roughness(cfg.roughness)
mat_inst.set_emissive(cfg.emissive)
# mat_inst.set_emissive_intensity(self.cfg.emissive_intensity) # Unimplemented
mat_inst.set_base_color_texture(cfg.base_color_texture)
mat_inst.set_metallic_texture(cfg.metallic_texture)
mat_inst.set_roughness_texture(cfg.roughness_texture)
mat_inst.set_normal_texture(cfg.normal_texture)
mat_inst.set_ao_texture(cfg.ao_texture)
if self.is_rt_enabled:
mat_inst.set_ior(cfg.ior)
mat_inst.mat.update_pbr_material_type(cfg.rt_material_type)
[docs]
def create_instance(self, uid: str) -> VisualMaterialInst:
"""Create a new material instance from this material template.
Note:
- If the uid already exists, the existing instance will be returned.
Args:
uid (str): Unique identifier for the material instance.
Returns:
VisualMaterialInst: The created material instance.
"""
inst = VisualMaterialInst(uid, self._mat)
# TODO: Support change default properties for material.
# This will improve the instance creation efficiency.
self.set_default_properties(inst, self.cfg)
return inst
[docs]
def get_default_instance(self) -> VisualMaterialInst:
"""Get the default material instance created with the same uid as the material template.
Returns:
VisualMaterialInst: The default material instance.
"""
return self._default_mat_inst
[docs]
def get_instance(self, uid: str) -> VisualMaterialInst:
"""Get an existing material instance by its uid.
Args:
uid (str): Unique identifier for the material instance.
Returns:
VisualMaterialInst: The material instance.
"""
return VisualMaterialInst(uid, self._mat)
[docs]
class VisualMaterialInst:
"""Instance of a visual material in the simulation environment."""
[docs]
def __init__(self, uid: str, mat: Material):
self.uid = uid
self._mat = mat
# Init properties with default values
self.base_color = [0.5, 0.5, 0.5, 1.0]
self.metallic = 0.0
self.roughness = 0.5
self.emissive = [0.0, 0.0, 0.0]
self.emissive_intensity = 1.0
self.base_color_texture = None
self.metallic_texture = None
self.roughness_texture = None
self.normal_texture = None
self.ao_texture = None
self.ior = 1.5
# self.subsurface = 0.0
@property
def mat(self) -> MaterialInst:
return self._mat.get_inst(self.uid)
[docs]
def set_base_color(self, color: list) -> None:
"""Set base color/diffuse color."""
self.base_color = color
self.mat.set_base_color(color)
[docs]
def set_roughness(self, roughness: float) -> None:
"""Set surface roughness."""
self.roughness = roughness
inst = self._mat.get_inst(self.uid)
inst.set_roughness(roughness)
[docs]
def set_emissive(self, emissive: list) -> None:
"""Set emissive color."""
self.emissive = emissive
value = np.zeros(4)
value[0:3] = emissive
inst = self._mat.get_inst(self.uid)
inst.set_emissive(value)
[docs]
def set_emissive_intensity(self, intensity: float) -> None:
"""Set emissive intensity multiplier."""
logger.log_error("Unimplemented: set_emissive_intensity")
[docs]
def set_base_color_texture(
self, texture_path: str = None, texture_data: Optional[torch.Tensor] = None
) -> None:
"""Set base color texture from file path or texture data.
Args:
texture_path: Path to texture file
texture_data: Texture data as a torch.Tensor
"""
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."
)
if texture_path is not None:
self.base_color_texture = texture_path
inst = self._mat.get_inst(self.uid)
inst.set_base_color_map(texture_path)
elif texture_data is not None:
self.base_color_texture = texture_data
inst = self._mat.get_inst(self.uid)
# TODO: Optimize texture creation method.
world = dexsim.default_world()
env = world.get_env()
color_texture = env.create_color_texture(
texture_data.cpu().numpy(), has_alpha=True
)
inst.set_base_color_map(color_texture)
[docs]
def set_metallic_texture(
self, texture_path: str = None, texture_data: Optional[torch.Tensor] = None
) -> None:
"""Set metallic texture from file path or texture data.
Args:
texture_path: Path to texture file
texture_data: Texture data as a torch.Tensor
"""
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."
)
if texture_path is not None:
self.metallic_texture = texture_path
inst = self._mat.get_inst(self.uid)
inst.set_metallic_map(texture_path)
elif texture_data is not None:
self.metallic_texture = texture_data
inst = self._mat.get_inst(self.uid)
# TODO: Optimize texture creation method.
world = dexsim.default_world()
env = world.get_env()
metallic_texture = env.create_color_texture(
texture_data.cpu().numpy(), has_alpha=False
)
inst.set_metallic_map(metallic_texture)
[docs]
def set_roughness_texture(
self, texture_path: str = None, texture_data: Optional[torch.Tensor] = None
) -> None:
"""Set roughness texture from file path or texture data.
Args:
texture_path: Path to texture file
texture_data: Texture data as a torch.Tensor
"""
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."
)
if texture_path is not None:
self.roughness_texture = texture_path
inst = self._mat.get_inst(self.uid)
inst.set_roughness_map(texture_path)
elif texture_data is not None:
self.roughness_texture = texture_data
inst = self._mat.get_inst(self.uid)
# TODO: Optimize texture creation method.
world = dexsim.default_world()
env = world.get_env()
roughness_texture = env.create_color_texture(
texture_data.cpu().numpy(), has_alpha=False
)
inst.set_roughness_map(roughness_texture)
[docs]
def set_normal_texture(
self, texture_path: str = None, texture_data: Optional[torch.Tensor] = None
) -> None:
"""Set normal texture from file path or texture data.
Args:
texture_path: Path to texture file
texture_data: Texture data as a torch.Tensor
"""
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."
)
if texture_path is not None:
self.normal_texture = texture_path
inst = self._mat.get_inst(self.uid)
inst.set_normal_map(texture_path)
elif texture_data is not None:
self.normal_texture = texture_data
inst = self._mat.get_inst(self.uid)
# TODO: Optimize texture creation method.
world = dexsim.default_world()
env = world.get_env()
normal_texture = env.create_color_texture(
texture_data.cpu().numpy(), has_alpha=False
)
inst.set_normal_map(normal_texture)
[docs]
def set_ao_texture(
self, texture_path: str = None, texture_data: Optional[torch.Tensor] = None
) -> None:
"""Set ambient occlusion texture from file path or texture data.
Args:
texture_path: Path to texture file
texture_data: Texture data as a torch.Tensor
"""
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."
)
if texture_path is not None:
self.ao_texture = texture_path
inst = self._mat.get_inst(self.uid)
inst.set_ao_map(texture_path)
elif texture_data is not None:
self.ao_texture = texture_data
inst = self._mat.get_inst(self.uid)
# TODO: Optimize texture creation method.
world = dexsim.default_world()
env = world.get_env()
ao_texture = env.create_color_texture(
texture_data.cpu().numpy(), has_alpha=False
)
inst.set_ao_map(ao_texture)
[docs]
def set_ior(self, ior: float) -> None:
"""Set index of refraction."""
if is_rt_enabled() is False:
logger.log_debug("Ray Tracing rendering not enabled, ignoring IOR setting.")
return
self.ior = ior
inst = self._mat.get_inst(self.uid)
inst.set_rt_param("ior", ior)