Implement basic USD test

PiperOrigin-RevId: 617302442
Change-Id: Ic598df9757288640b5f35c4e740958884b1b4698

python/mujoco/testdata/usd_golden.usda

@@ -0,0 +1,71 @@
+#usda 1.0
+(
+    defaultPrim = "World"
+    doc = """Generated from Composed Stage of root layer 
+"""
+    endTimeCode = 1
+    startTimeCode = 0
+    timeCodesPerSecond = 24
+    upAxis = "Z"
+)
+
+def Xform "World"
+{
+    def Xform "Light_Xform_0"
+    {
+        double3 xformOp:translate
+        uniform token[] xformOpOrder = ["xformOp:translate"]
+
+        def SphereLight "Light_0"
+        {
+            bool inputs:normalize = 1
+            float inputs:radius = 0.3
+            bool treatAsPoint = 0
+        }
+    }
+
+    def Xform "CubeMesh_Xform_white_box"
+    {
+        matrix4d xformOp:transform.timeSamples = {
+            0: ( (0, 0, 0, 0), (0, 0, 0, 0), (0, 0, 0, 0), (0, 0, 0, 1) ),
+        }
+        uniform token[] xformOpOrder = ["xformOp:transform"]
+
+        def Mesh "CubeMesh_white_box" (
+            apiSchemas = ["MaterialBindingAPI"]
+        )
+        {
+            int[] faceVertexCounts = [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]
+            int[] faceVertexIndices = [4, 7, 5, 4, 6, 7, 0, 2, 4, 2, 6, 4, 0, 1, 2, 1, 3, 2, 1, 5, 7, 1, 7, 3, 2, 3, 7, 2, 7, 6, 0, 4, 1, 1, 4, 5]
+            rel material:binding = </World/_materials/Material_white_box>
+            point3f[] points = [(-1, -1, -1), (1, -1, -1), (-1, -1, 1), (1, -1, 1), (-1, 1, -1), (1, 1, -1), (-1, 1, 1), (1, 1, 1)]
+            texCoord2f[] primvars:UVMap = [] (
+                interpolation = "faceVarying"
+            )
+            int[] primvars:UVMap:indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35]
+            uniform token subdivisionScheme = "none"
+            token visibility.timeSamples = {
+                0: "inherited",
+            }
+        }
+    }
+
+    def "_materials"
+    {
+        def Material "Material_white_box"
+        {
+            token outputs:surface.connect = </World/_materials/Material_white_box/Principled_BSDF.outputs:surface>
+
+            def Shader "Principled_BSDF"
+            {
+                uniform token info:id = "UsdPreviewSurface"
+                color3f inputs:diffuseColor = (1, 1, 1)
+                float inputs:metallic = 0.5
+                float inputs:opacity = 1
+                float inputs:roughness = 0.5
+                token outputs:surface
+            }
+        }
+    }
+}
+

python/mujoco/usd/component.py

@@ -15,6 +15,7 @@
 from typing import List, Optional, Tuple
 
 import mujoco
+import mujoco.usd.utils
 import numpy as np
 
 # TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
@@ -36,9 +37,9 @@ class USDMesh:
       stage: Usd.Stage,
       model: mujoco.MjModel,
       geom: mujoco.MjvGeom,
-      objid: int,
+      objid: str,
       dataid: int,
-      rgba: Tuple[int, ...] = (1, 1, 1, 1),
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
       texture_file: Optional[str] = None,
   ):
     self.stage = stage
@@ -216,7 +217,7 @@ class USDMesh:
     UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
 
   def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
-    transformation_mat = mujoco.usd_utils.create_transform_matrix(
+    transformation_mat = mujoco.usd.utils.create_transform_matrix(
         rotation_matrix=mat, translation_vector=pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
@@ -235,8 +236,8 @@ class USDPrimitiveMesh:
       self,
       stage: Usd.Stage,
       geom: mujoco.MjvGeom,
-      objid: int,
+      objid: str,
-      rgba: Tuple[int, ...] = (1, 1, 1, 1),
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
       texture_file: Optional[str] = None,
   ):
     self.stage = stage
@@ -246,9 +247,9 @@ class USDPrimitiveMesh:
     self.texture_file = texture_file
 
     self.usd_prim = Usd.Prim()
-    self.usd_mesh = Usd.Mesh()
+    self.usd_mesh = UsdGeom.Mesh()
-    self.prim_mesh = Usd.PrimMesh()
+    self.prim_mesh = None
-    self.transform_op = Usd.TransformOp()
+    self.transform_op = Gf.Matrix4d(1.)
 
   def _set_refinement_properties(self):
     self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
@@ -364,7 +365,7 @@ class USDPrimitiveMesh:
     UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
 
   def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
-    transformation_mat = mujoco.usd_util.create_transform_matrix(
+    transformation_mat = mujoco.usd.utils.create_transform_matrix(
         rotation_matrix=mat, translation_vector=pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
@@ -383,8 +384,8 @@ class USDPrimitive:
       self,
       stage: Usd.Stage,
       geom: mujoco.MjvGeom,
-      objid: int,
+      objid: str,
-      rgba: Tuple[int, ...] = (1, 1, 1, 1),
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
       texture_file: Optional[str] = None,
   ):
     self.stage = stage
@@ -395,7 +396,7 @@ class USDPrimitive:
 
     self.usd_prim = Usd.Prim()
     self.usd_primitive_shape = Usd.PrimitiveShape()
-    self.transform_op = Usd.TransformOp()
+    self.transform_op = Gf.Matrix4d(1.)
 
   def _set_refinement_properties(self):
     self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
@@ -481,7 +482,7 @@ class USDPrimitive:
     UsdShade.MaterialBindingAPI(self.usd_primitive_shape).Bind(mtl)
 
   def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
-    transformation_mat = mujoco.usd_util.create_transform_matrix(
+    transformation_mat = mujoco.usd.utils.create_transform_matrix(
         rotation_matrix=mat, translation_vector=pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
@@ -500,8 +501,8 @@ class USDCapsule(USDPrimitive):
       self,
       stage: Usd.Stage,
       geom: mujoco.MjvGeom,
-      objid: int,
+      objid: str,
-      rgba: Tuple[int, ...] = (1, 1, 1, 1),
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
       texture_file: Optional[str] = None,
   ):
 
@@ -536,8 +537,8 @@ class USDEllipsoid(USDPrimitive):
       self,
       stage: Usd.Stage,
       geom: mujoco.MjvGeom,
-      objid: int,
+      objid: str,
-      rgba: Tuple[int, ...] = (1, 1, 1, 1),
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
       texture_file: Optional[str] = None,
   ):
 
@@ -569,8 +570,8 @@ class USDCubeMesh(USDPrimitiveMesh):
       self,
       stage: Usd.Stage,
       geom: mujoco.MjvGeom,
-      objid: int,
+      objid: str,
-      rgba: Tuple[int, ...] = (1, 1, 1, 1),
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
       texture_file: Optional[str] = None,
   ):
 
@@ -586,8 +587,6 @@ class USDCubeMesh(USDPrimitiveMesh):
         width=self.geom.size[0] * 2,
         height=self.geom.size[1] * 2,
         depth=self.geom.size[2] * 2,
-        create_uv_map=True,
-        map_texture_to_each_face=True,
     )
 
     self.prim_mesh.translate(-self.prim_mesh.get_center())
@@ -623,8 +622,8 @@ class USDSphereMesh(USDPrimitiveMesh):
       self,
       stage: Usd.Stage,
       geom: mujoco.MjvGeom,
-      objid: int,
+      objid: str,
-      rgba: Tuple[int, ...] = (1, 1, 1, 1),
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
       texture_file: Optional[str] = None,
   ):
 
@@ -637,7 +636,7 @@ class USDSphereMesh(USDPrimitiveMesh):
     self.usd_prim = stage.GetPrimAtPath(mesh_path)
 
     self.prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
-        radius=float(self.geom.size[0]), create_uv_map=True
+        radius=float(self.geom.size[0])
     )
 
     self.prim_mesh.translate(-self.prim_mesh.get_center())
@@ -673,8 +672,8 @@ class USDCylinderMesh(USDPrimitiveMesh):
       self,
       stage: Usd.Stage,
       geom: mujoco.MjvGeom,
-      objid: int,
+      objid: str,
-      rgba: Tuple[int, ...] = (1, 1, 1, 1),
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
       texture_file: Optional[str] = None,
   ):
 
@@ -689,7 +688,6 @@ class USDCylinderMesh(USDPrimitiveMesh):
     self.prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
         radius=self.geom.size[0],
         height=self.geom.size[2] * 2,
-        create_uv_map=True,
     )
 
     self.prim_mesh.translate(-self.prim_mesh.get_center())
@@ -725,8 +723,8 @@ class USDPlaneMesh(USDPrimitiveMesh):
       self,
       stage: Usd.Stage,
       geom: mujoco.MjvGeom,
-      objid: int,
+      objid: str,
-      rgba: Tuple[int, ...] = (1, 1, 1, 1),
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
       texture_file: Optional[str] = None,
   ):
 
@@ -742,8 +740,6 @@ class USDPlaneMesh(USDPrimitiveMesh):
         width=self.geom.size[0] * 2 if self.geom.size[0] > 0 else 100,
         height=self.geom.size[1] * 2 if self.geom.size[1] > 0 else 100,
         depth=0.001,
-        create_uv_map=True,
-        map_texture_to_each_face=True,
     )
 
     self.prim_mesh.translate(-self.prim_mesh.get_center())
@@ -776,7 +772,7 @@ class USDPlaneMesh(USDPrimitiveMesh):
 class USDSphereLight:
 
   def __init__(
-      self, stage: Usd.Stage, objid: int, radius: Optional[float] = 0.3
+      self, stage: Usd.Stage, objid: str, radius: Optional[float] = 0.3
   ):
     self.stage = stage
 
@@ -808,7 +804,7 @@ class USDSphereLight:
 
 class USDDomeLight:
 
-  def __init__(self, stage: Usd.Stage, objid: int):
+  def __init__(self, stage: Usd.Stage, objid: str):
     self.stage = stage
 
     xform_path = f"/World/Light_Xform_{objid}"
@@ -829,7 +825,7 @@ class USDDomeLight:
 
 class USDCamera:
 
-  def __init__(self, stage: Usd.Stage, objid: int):
+  def __init__(self, stage: Usd.Stage, objid: str):
     self.stage = stage
 
     xform_path = f"/World/Camera_Xform_{objid}"
@@ -851,7 +847,7 @@ class USDCamera:
 
   def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
 
-    transformation_mat = mujoco.usd_util.create_transform_matrix(
+    transformation_mat = mujoco.usd.utils.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

@@ -15,6 +15,7 @@
 import os
 
 import mujoco
+import mujoco.usd.component as component_module
 import numpy as np
 import scipy
 import termcolor
@@ -235,7 +236,7 @@ class USDExporter:
 
     # handles meshes in scene
     if geom.type == mujoco.mjtGeom.mjGEOM_MESH:
-      usd_geom = mujoco.usd.USDMesh(
+      usd_geom = component_module.USDMesh(
           stage=self.stage,
           model=self.model,
           geom=geom,
@@ -245,7 +246,7 @@ class USDExporter:
           texture_file=texture_file,
       )
     elif geom.type == mujoco.mjtGeom.mjGEOM_PLANE:
-      usd_geom = mujoco.usd.USDPlaneMesh(
+      usd_geom = component_module.USDPlaneMesh(
           stage=self.stage,
           geom=geom,
           objid=geom_name,
@@ -253,7 +254,7 @@ class USDExporter:
           texture_file=texture_file,
       )
     elif geom.type == mujoco.mjtGeom.mjGEOM_SPHERE:
-      usd_geom = mujoco.usd.USDSphereMesh(
+      usd_geom = component_module.USDSphereMesh(
           stage=self.stage,
           geom=geom,
           objid=geom_name,
@@ -261,7 +262,7 @@ class USDExporter:
           texture_file=texture_file,
       )
     elif geom.type == mujoco.mjtGeom.mjGEOM_CAPSULE:
-      usd_geom = mujoco.usd.USDCapsule(
+      usd_geom = component_module.USDCapsule(
           stage=self.stage,
           geom=geom,
           objid=geom_name,
@@ -269,7 +270,7 @@ class USDExporter:
           texture_file=texture_file,
       )
     elif geom.type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
-      usd_geom = mujoco.usd.USDEllipsoid(
+      usd_geom = component_module.USDEllipsoid(
           stage=self.stage,
           geom=geom,
           objid=geom_name,
@@ -277,7 +278,7 @@ class USDExporter:
           texture_file=texture_file,
       )
     elif geom.type == mujoco.mjtGeom.mjGEOM_CYLINDER:
-      usd_geom = mujoco.usd.USDCylinderMesh(
+      usd_geom = component_module.USDCylinderMesh(
           stage=self.stage,
           geom=geom,
           objid=geom_name,
@@ -285,7 +286,7 @@ class USDExporter:
           texture_file=texture_file,
       )
     elif geom.type == mujoco.mjtGeom.mjGEOM_BOX:
-      usd_geom = mujoco.usd.USDCubeMesh(
+      usd_geom = component_module.USDCubeMesh(
           stage=self.stage,
           geom=geom,
           objid=geom_name,
@@ -331,7 +332,8 @@ class USDExporter:
     for i in range(self.scene.nlight):
       light = self.scene.lights[i]
       if not np.allclose(light.pos, [0, 0, 0]):
-        self.usd_lights.append(mujoco.usd.USDSphereLight(stage=self.stage, objid=i))
+        self.usd_lights.append
+        (component_module.USDSphereLight(stage=self.stage, objid=str(i)))
       else:
         self.usd_lights.append(None)
 
@@ -342,7 +344,7 @@ class USDExporter:
       if np.allclose(light.pos, [0, 0, 0]):
         continue
 
-      if self.usd_lights[i] is None:
+      if i >= len(self.usd_lights) or self.usd_lights[i] is None:
         continue
 
       self.usd_lights[i].update(
@@ -356,8 +358,10 @@ class USDExporter:
 
   def _load_cameras(self):
     self.usd_cameras = []
-    for name in self.camera_names:
+    if self.camera_names is not None:
-      self.usd_cameras.append(mujoco.usd.USDCamera(stage=self.stage, objid=name))
+      for name in self.camera_names:
+        self.usd_cameras.append(
+            component_module.USDCamera(stage=self.stage, objid=name))
 
   def _update_cameras(
       self,
@@ -397,11 +401,11 @@ class USDExporter:
   ):
 
     if light_type == "sphere":
-      new_light = mujoco.usd.USDSphereLight(stage=self.stage, objid=objid, radius=radius)
+      new_light = component_module.USDSphereLight(stage=self.stage, objid=str(objid), radius=radius)
 
-      new_light.update(pos=pos, intensity=intensity, color=color, frame=0)
+      new_light.update(pos=np.array(pos), intensity=intensity, color=color, frame=0)
     elif light_type == "dome":
-      new_light = mujoco.usd.USDDomeLight(stage=self.stage, objid=objid)
+      new_light = component_module.USDDomeLight(stage=self.stage, objid=str(objid))
 
       new_light.update(intensity=intensity, color=color, frame=0)
 
@@ -411,11 +415,11 @@ class USDExporter:
       rotation_xyz: List[float],
       objid: Optional[int] = 1,
   ):
-    new_camera = mujoco.usd.USDCamera(stage=self.stage, objid=objid)
+    new_camera = component_module.USDCamera(stage=self.stage, objid=str(objid))
 
     r = scipy.spatial.transform.Rotation.from_euler(
         "xyz", rotation_xyz, degrees=True)
-    new_camera.update(cam_pos=pos, cam_mat=r.as_matrix(), frame=0)
+    new_camera.update(cam_pos=np.array(pos), cam_mat=r.as_matrix(), frame=0)
 
   def save_scene(self, filetype: str = "usd"):
     self.stage.SetEndTimeCode(self.frame_count)

python/mujoco/usd/exporter_test.py

@@ -0,0 +1,70 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# 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.
+# ==============================================================================
+"""Tests for the MuJoCo USD Exporter."""
+
+import logging
+import os
+
+import mujoco
+
+from absl.testing import absltest
+from etils import epath
+
+# Open3D and USD are not fully supported on all MuJoCo architectures.
+# pylint: disable=python.style(g-import-not-at-top)
+execute_test = True
+try:
+  from mujoco.usd import exporter as exporter_module
+  from pxr import Usd
+except ImportError:
+  logging.warning('Skipping test due to missing import')
+  execute_test = False
+# pylint: enable=python.style(g-import-not-at-top)
+
+class ExporterTest(absltest.TestCase):
+
+  def test_usd_export(self):
+    if not execute_test:
+      return
+
+    output_dir = os.getenv('TEST_UNDECLARED_OUTPUTS_DIR')
+    xml = """
+<mujoco>
+  <worldbody>
+    <camera name="closeup" pos="0 -6 0" xyaxes="1 0 0 0 1 100"/>
+    <geom name="white_box" type="box" size="1 1 1" rgba="1 1 1 1"/>
+  </worldbody>
+</mujoco>
+"""
+    model = mujoco.MjModel.from_xml_string(xml)
+    data = mujoco.MjData(model)
+    exporter = exporter_module.USDExporter(
+        model,
+        output_directory_name="mujoco_usdpkg",
+        output_directory_root=output_dir,
+    )
+    exporter.update_scene(data)
+    exporter.save_scene("export.usda")
+
+    with open(os.path.join(
+        output_dir, "mujoco_usdpkg/frames", "frame_1_.export.usda"), "r") as f:
+      golden_path = os.path.join(
+          epath.resource_path("mujoco"), "testdata", "usd_golden.usda")
+      with open(golden_path, "r") as golden_file:
+        self.assertEqual(f.read(), golden_file.read())
+
+
+if __name__ == "__main__":
+  absltest.main()