[FEATURE] Add USD Materials Baking.

Author: duburcqa

.github/workflows/generic.yml

@@ -32,6 +32,8 @@ jobs:
       TI_ENABLE_METAL: "0"
       TI_ENABLE_OPENGL: "0"
       TI_ENABLE_VULKAN: "0"
+      OMNI_KIT_ACCEPT_EULA: "yes"
+      OMNI_KIT_ALLOW_ROOT: "1"
 
     runs-on: ${{ matrix.OS }}
     if: github.event_name != 'release'
@@ -108,7 +110,13 @@ jobs:
 
       - name: Install Genesis
         run: |
-          pip install -e '.[dev]'
+          # TODO: USD baking does not currently support Python 3.11 since
+          # NVIDIA does not currently release `omniverse-kit==107.3` on PyPI.
+          # See: https://github.com/Genesis-Embodied-AI/Genesis/pull/1300
+          pip install -e '.[dev,usd]'
+          if [[ "${{ matrix.PYTHON_VERSION }}" == '3.10' ]] ; then
+            pip install --extra-index-url https://pypi.nvidia.com/ omniverse-kit
+          fi
 
       - name: Run unit tests
         run: |

.github/workflows/production.yml

@@ -25,6 +25,8 @@ jobs:
       GENESIS_IMAGE_VER: "1_1"
       TIMEOUT_MINUTES: 180
       MADRONA_DISABLE_CUDA_HEAP_SIZE: "1"
+      OMNI_KIT_ACCEPT_EULA: "yes"
+      OMNI_KIT_ALLOW_ROOT: "1"
 
     steps:
       - name: Checkout code
@@ -60,7 +62,8 @@ jobs:
             --partition=hpc-mid ${NODELIST} --nodes=1 --time="${TIMEOUT_MINUTES}" \
             --job-name=${SLURM_JOB_NAME} \
             bash -c "
-              pip install -e '.[dev,render]' && \
+              pip install --extra-index-url https://pypi.nvidia.com/ omniverse-kit && \
+              pip install -e '.[dev,render,usd]' && \
               pytest -v --forked ./tests
             "
 

genesis/assets/tests/chopper.glb

@@ -13,7 +13,6 @@
 from genesis.options.surfaces import Surface
 import genesis.utils.mesh as mu
 import genesis.utils.gltf as gltf_utils
-import genesis.utils.usda as usda_utils
 import genesis.utils.particle as pu
 from genesis.repr_base import RBC
 
@@ -229,6 +228,7 @@ def from_trimesh(
         """
         if surface is None:
             surface = gs.surfaces.Default()
+            surface.update_texture()
         else:
             surface = surface.copy()
         mesh = mesh.copy(include_cache=True)
@@ -341,22 +341,20 @@ def from_morph_surface(cls, morph, surface=None):
         If the morph is a Mesh morph (morphs.Mesh), it could contain multiple submeshes, so we return a list.
         """
         if isinstance(morph, gs.options.morphs.Mesh):
-            if morph.file.endswith(("obj", "ply", "stl")):
+            if morph.is_format(gs.options.morphs.MESH_FORMATS):
                 meshes = mu.parse_mesh_trimesh(morph.file, morph.group_by_material, morph.scale, surface)
-
-            elif morph.file.endswith(("glb", "gltf")):
+            elif morph.is_format(gs.options.morphs.GLTF_FORMATS):
                 if morph.parse_glb_with_trimesh:
                     meshes = mu.parse_mesh_trimesh(morph.file, morph.group_by_material, morph.scale, surface)
                 else:
                     meshes = gltf_utils.parse_mesh_glb(morph.file, morph.group_by_material, morph.scale, surface)
+            elif morph.is_format(gs.options.morphs.USD_FORMATS):
+                import genesis.utils.usda as usda_utils
 
-            elif morph.file.endswith(("usd", "usda", "usdc", "usdz")):
                 meshes = usda_utils.parse_mesh_usd(morph.file, morph.group_by_material, morph.scale, surface)
-
             elif isinstance(morph, gs.options.morphs.MeshSet):
                 assert all(isinstance(mesh, trimesh.Trimesh) for mesh in morph.files)
                 meshes = [mu.trimesh_to_mesh(mesh, morph.scale, surface) for mesh in morph.files]
-
             else:
                 gs.raise_exception(
                     f"File type not supported (yet). Submit a feature request if you need this: {morph.file}."

genesis/assets/tests/combined_srt.glb

@@ -1,3 +1,10 @@
+"""
+We define all types of morphologies here: shape primitives, meshes, URDF, MJCF, and soft robot description files.
+
+These are independent of backend solver type and are shared by different solvers, e.g. a mesh can be either loaded as a
+rigid object / MPM object / FEM object.
+"""
+
 import os
 from typing import Any, List, Optional, Sequence, Tuple, Union
 
@@ -10,11 +17,12 @@
 from .misc import CoacdOptions
 from .options import Options
 
-"""
-We define all types of morphologies here: shape primitives, meshes, URDF, MJCF, and soft robot description files.
-These are independent of backend solver type and are shared by different solvers.
-E.g. a mesh can be either loaded as a rigid object / MPM object / FEM object.
-"""
+
+URDF_FORMAT = ".urdf"
+MJCF_FORMAT = ".xml"
+MESH_FORMATS = (".obj", ".ply", ".stl")
+GLTF_FORMATS = (".glb", ".gltf")
+USD_FORMATS = (".usd", ".usda", ".usdc", ".usdz")
 
 
 class TetGenMixin(Options):
@@ -556,6 +564,9 @@ def __init__(self, **data):
     def _repr_type(self):
         return f"<gs.morphs.{self.__class__.__name__}(file='{self.file}')>"
 
+    def is_format(self, format):
+        return self.file.lower().endswith(format)
+
 
 class Mesh(FileMorph, TetGenMixin):
     """
@@ -768,8 +779,8 @@ class MJCF(FileMorph):
 
     def __init__(self, **data):
         super().__init__(**data)
-        if not self.file.endswith(".xml"):
-            gs.raise_exception(f"Expected `.xml` extension for MJCF file: {self.file}")
+        if not self.is_format(MJCF_FORMAT):
+            gs.raise_exception(f"Expected `{MJCF_FORMAT}` extension for MJCF file: {self.file}")
 
         # What you want to do with scaling is kinda "zoom" the world from the perspective of the entity, i.e. scale the
         # geometric properties of an entity wrt its root pose. In the general case, ie for a 3D vector scale, (x, y, z)
@@ -878,8 +889,8 @@ class URDF(FileMorph):
 
     def __init__(self, **data):
         super().__init__(**data)
-        if isinstance(self.file, str) and not self.file.endswith(".urdf"):
-            gs.raise_exception(f"Expected `.urdf` extension for URDF file: {self.file}")
+        if isinstance(self.file, str) and not self.is_format(URDF_FORMAT):
+            gs.raise_exception(f"Expected `{URDF_FORMAT}` extension for URDF file: {self.file}")
 
         # Anisotropic scaling is ill-defined for poly-articulated robots. See related MJCF about this for details.
         if isinstance(self.scale, np.ndarray) and self.scale.std() > gs.EPS:
@@ -997,10 +1008,10 @@ def __init__(self, **data):
         # Make sure that Propellers links are preserved
         self.links_to_keep = tuple(set([*self.links_to_keep, *self.propellers_link_name]))
 
-        if isinstance(self.file, str) and not self.file.endswith(".urdf"):
-            gs.raise_exception(f"Drone only supports `.urdf` extension: {self.file}")
+        if isinstance(self.file, str) and not self.is_format(URDF_FORMAT):
+            gs.raise_exception(f"Drone only supports `{URDF_FORMAT}` extension: {self.file}")
 
-        if self.model not in ["CF2X", "CF2P", "RACE"]:
+        if self.model not in ("CF2X", "CF2P", "RACE"):
             gs.raise_exception(f"Unsupported `model`: {self.model}.")
 
 

genesis/assets/tests/combined_transform.glb

@@ -97,6 +97,9 @@ def __init__(self, **data):
             else:
                 gs.logger.warning("`env_euler` is ignored when `env_quat` is specified.")
 
+        if self.env_surface is not None:
+            self.env_surface.update_texture()
+
 
 class BatchRenderer(RendererOptions):
     """

genesis/engine/mesh.py

@@ -7,6 +7,7 @@
 from PIL import Image
 
 import genesis as gs
+import genesis.utils.mesh as mu
 
 from .options import Options
 
@@ -32,6 +33,9 @@ def check_simplify(self):
     def apply_cutoff(self, cutoff):
         raise NotImplementedError
 
+    def is_black(self):
+        raise NotImplementedError
+
 
 class ColorTexture(Texture):
     """
@@ -67,6 +71,9 @@ def apply_cutoff(self, cutoff):
             return
         self.color = tuple(1.0 if c >= cutoff else 0.0 for c in self.color)
 
+    def is_black(self):
+        return all(c < gs.EPS for c in self.color)
+
 
 class ImageTexture(Texture):
     """
@@ -115,25 +122,7 @@ def __init__(self, **data):
             if self.image_path.endswith((".hdr", ".exr")):
                 self.encoding = "linear"  # .exr or .hdr images should be encoded with 'linear'
                 if self.image_path.endswith((".exr")):
-                    exr_file = OpenEXR.InputFile(self.image_path)
-                    exr_header = exr_file.header()
-
-                    if exr_header["compression"].v > Imath.Compression.PIZ_COMPRESSION:
-                        new_image_path = f"{self.image_path[:-4]}_ZIP.exr"
-                        gs.logger.warning(
-                            f"EXR image {self.image_path}'s compression type {exr_header['compression']} is not supported. "
-                            f"Converting to compression type ZIP_COMPRESSION and saving to {new_image_path}."
-                        )
-                        self.image_path = new_image_path
-
-                        if not os.path.exists(new_image_path):
-                            channel_data = {channel: exr_file.channel(channel) for channel in exr_header["channels"]}
-                            exr_header["compression"] = Imath.Compression(Imath.Compression.ZIP_COMPRESSION)
-                            new_exr_file = OpenEXR.OutputFile(new_image_path, exr_header)
-                            new_exr_file.writePixels(channel_data)
-                            new_exr_file.close()
-
-                    exr_file.close()
+                    self.image_path = mu.check_exr_compression(self.image_path)
             else:
                 self.image_array = np.array(Image.open(self.image_path))
                 self.image_path = None
@@ -158,16 +147,8 @@ def __init__(self, **data):
                     )
             self.image_array = arr
 
-        # calculate channel
-        if self.image_array is None:
-            if isinstance(self.resolution, (tuple, list)):
-                H, W = self.resolution
-            else:
-                H = W = self.resolution
-
-            # Default to 3-channel RGB
-            white = np.array([255, 255, 255], dtype=np.uint8)
-            self.image_array = np.full((H, W, 3), white, dtype=np.uint8)
+        # just calculate channel
+        if self.image_array is None:  # Using 'image_path'
             self._mean_color = np.array([1.0, 1.0, 1.0], dtype=np.float16)
             self._channel = 3
         else:
@@ -216,3 +197,6 @@ def apply_cutoff(self, cutoff):
         if cutoff is None or self.image_array is None:  # Cutoff does not apply on image file.
             return
         self.image_array = np.where(self.image_array >= 255.0 * cutoff, 255, 0).astype(np.uint8)
+
+    def is_black(self):
+        return all(c < gs.EPS for c in self.image_color) or np.max(self.image_array) == 0

genesis/options/morphs.py

@@ -219,8 +219,9 @@ def parse_glb_material(glb, material_index, surface):
         if material.emissiveFactor is not None:
             emissive_factor = np.array(material.emissiveFactor, dtype=np.float32)
 
-        if emissive_factor is not None and np.any(emissive_factor > 0.0):  # Make sure to check emissive
-            emissive_texture = mu.create_texture(emissive_image, emissive_factor, "srgb")
+        emissive_texture = mu.create_texture(emissive_image, emissive_factor, "srgb")
+        if emissive_texture.is_black():  # Make sure to check emissive
+            emissive_texture = None
 
     # TODO: Parse them!
     for extension_name, extension_material in material.extensions.items():
@@ -317,7 +318,7 @@ def parse_mesh_glb(path, group_by_material, scale, surface):
                         primitive.material, parse_glb_material(glb, primitive.material, surface)
                     )
             else:
-                material, uv_used, material_name = None, 0, ""
+                material, uv_used, material_name = surface.copy(), 0, ""
 
             uvs = None
             if "KHR_draco_mesh_compression" in primitive.extensions: