Updating for PR comments

python/mujoco/usd/camera.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 """Camera handling for USD exporter."""
 
-import mujoco.usd.utils as utils_modules
+import mujoco.usd.utils as utils_module
 
 import numpy as np
 
@@ -49,7 +49,7 @@ class USDCamera:
 
   def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
     """Updates the position and orientation of the camera in the scene."""
-    transformation_mat = utils_modules.create_transform_matrix(
+    transformation_mat = utils_module.create_transform_matrix(
         rotation_matrix=cam_mat, translation_vector=cam_pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)

python/mujoco/usd/exporter.py

@@ -251,11 +251,10 @@ class USDExporter:
 
     assert geom_name not in self.geom_names
 
-    geom_textures = (
-      [(self.texture_files[i], self.model.tex_type[i]) if i != -1 else None for i in self.model.mat_texid[geom.matid]]
-      if geom.matid != -1
-      else None
-    )
+    if geom.matid == -1:
+      geom_textures = []
+    else:
+      geom_textures = [(self.texture_files[i], self.model.tex_type[i]) if i != -1 else None for i in self.model.mat_texid[geom.matid]]
 
     # handling meshes in our scene
     if geom.type == mujoco.mjtGeom.mjGEOM_MESH:

python/mujoco/usd/objects.py

@@ -16,7 +16,7 @@
 
 import abc
 import collections
-from typing import Optional, Dict, Any, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import mujoco
 import mujoco.usd.shapes as shapes_module
@@ -55,7 +55,7 @@ class USDObject(abc.ABC):
       geom: mujoco.MjvGeom,
       obj_name: str,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
+      geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
   ):
     self.stage = stage
     self.model = model
@@ -225,7 +225,7 @@ class USDMesh(USDObject):
       obj_name: str,
       dataid: int,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
+      geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
   ):
     super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
 
@@ -243,7 +243,7 @@ class USDMesh(USDObject):
     )
     self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
 
-    if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
+    if geom.matid != -1 and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
       # setting mesh uv properties
       mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
       self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
@@ -319,7 +319,7 @@ class USDPrimitiveMesh(USDObject):
       geom: mujoco.MjvGeom,
       obj_name: str,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
+      geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
   ):
     super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
 
@@ -338,7 +338,7 @@ class USDPrimitiveMesh(USDObject):
     self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
     self._set_refinement_properties(self.usd_prim)
 
-    if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
+    if geom.matid != -1 and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
       # setting mesh uv properties
       mesh_texcoord, _ = self._get_uv_geometry()
       self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
@@ -403,7 +403,7 @@ class USDTendon(USDObject):
       geom: mujoco.MjvGeom,
       obj_name: str,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
+      geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
   ):
     super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
 
@@ -435,7 +435,7 @@ class USDTendon(USDObject):
           part_geometry["mesh_face"]
       )
 
-    if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
+    if geom.matid != -1 and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
       # setting uv properties for each of the parts in the tendon
       part_uv_geometries = self._get_uv_geometry()
       for name, part_uv_geometry in part_uv_geometries.items():

python/mujoco/usd/shapes.py

@@ -14,14 +14,14 @@
 # ==============================================================================
 """Built-in shapes for USD exporter."""
 
-from typing import Dict, Any, Tuple, Optional, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import mujoco
 import numpy as np
 
 def get_triangle_uvs(
-    vertices: np.array,
-    triangles: np.array,
+    vertices: np.ndarray,
+    triangles: np.ndarray,
     texture_type: Optional[mujoco.mjtTexture]
 ):
   if texture_type == None:
@@ -70,13 +70,12 @@ def get_triangle_uvs(
 
   return np.array(triangle_uvs)
 
-class TriangleMesh():
-  """ Store UV and geometry information for a primitve mesh
-  """
+class TriangleMesh:
+  """Store UV and geometry information for a primitive mesh."""
   def __init__(self,
-               vertices: np.array,
-               triangles: np.array,
-               triangle_uvs: np.array):
+               vertices: np.ndarray,
+               triangles: np.ndarray,
+               triangle_uvs: np.ndarray):
     self.vertices = vertices
     self.triangles = triangles
     self.triangle_uvs = triangle_uvs
@@ -88,7 +87,7 @@ class TriangleMesh():
       height: float,
       depth: float,
       texture_type: Optional[mujoco.mjtTexture]
-  ):
+  ) -> 'TriangleMesh':
     vertices = np.array([[0.0, 0.0, 0.0],
                          [width, 0.0, 0.0],
                          [0.0, 0.0, depth],
@@ -121,7 +120,7 @@ class TriangleMesh():
       radius: float,
       texture_type: Optional[mujoco.mjtTexture],
       resolution: int
-  ):
+  ) -> 'TriangleMesh':
     vertices = []
     triangles = []
     for i in range(2*resolution + 1):
@@ -154,7 +153,7 @@ class TriangleMesh():
     radius: float,
     texture_type: Optional[mujoco.mjtTexture],
     resolution: int
-  ):
+  ) -> 'TriangleMesh':
     vertices = []
     triangles = []
     for i in range(resolution + 1):
@@ -193,7 +192,7 @@ class TriangleMesh():
       height: float,
       texture_type: Optional[mujoco.mjtTexture],
       resolution: int
-  ):
+  ) -> 'TriangleMesh':
     vertices = []
     triangles = []