Merge pull request #1901 from abhihjoshi:usd-fixes

PiperOrigin-RevId: 671815757
Change-Id: I8a08385b63401c743253545d89389ab8aaac91bc

.github/workflows/build.yml

@@ -142,6 +142,7 @@ jobs:
         fi
         source venv/bin/activate
         python -m pip install --upgrade --require-hashes -r "${repo}/python/build_requirements.txt"
+        python -m pip install --upgrade --require-hashes -r "${repo}/python/build_requirements_usd.txt"
     - name: Configure MuJoCo
       run: >
         mkdir build &&

python/build_requirements_usd.txt

@@ -0,0 +1,41 @@
+# usd-core is not available for Linux aarch64, so leave it at the end incase it
+# fails.
+pillow==10.2.0 \
+    --hash=sha256:a086c2af425c5f62a65e12fbf385f7c9fcb8f107d0849dba5839461a129cf311 \
+    --hash=sha256:35bb52c37f256f662abdfa49d2dfa6ce5d93281d323a9af377a120e89a9eafb5 \
+    --hash=sha256:1da3b2703afd040cf65ec97efea81cfba59cdbed9c11d8efc5ab09df9509fc56 \
+    --hash=sha256:0eae2073305f451d8ecacb5474997c08569fb4eb4ac231ffa4ad7d342fdc25ac \
+    --hash=sha256:1e23412b5c41e58cec602f1135c57dfcf15482013ce6e5f093a86db69646a5aa \
+    --hash=sha256:154e939c5f0053a383de4fd3d3da48d9427a7e985f58af8e94d0b3c9fcfcf4f9 \
+    --hash=sha256:9c23f307202661071d94b5e384e1e1dc7dfb972a28a2310e4ee16103e66ddb67 \
+    --hash=sha256:11fa2e5984b949b0dd6d7a94d967743d87c577ff0b83392f17cb3990d0d2fd6e \
+    --hash=sha256:83b2021f2ade7d1ed556bc50a399127d7fb245e725aa0113ebd05cfe88aaf588 \
+    --hash=sha256:c570f24be1e468e3f0ce7ef56a89a60f0e05b30a3669a459e419c6eac2c35364 \
+    --hash=sha256:da2b52b37dad6d9ec64e653637a096905b258d2fc2b984c41ae7d08b938a67e4 \
+    --hash=sha256:aebb6044806f2e16ecc07b2a2637ee1ef67a11840a66752751714a0d924adf72 \
+    --hash=sha256:870ea1ada0899fd0b79643990809323b389d4d1d46c192f97342eeb6ee0b8483 \
+    --hash=sha256:c365fd1703040de1ec284b176d6af5abe21b427cb3a5ff68e0759e1e313a5e7e \
+    --hash=sha256:0304004f8067386b477d20a518b50f3fa658a28d44e4116970abfcd94fac34a8 \
+    --hash=sha256:1b5e1b74d1bd1b78bc3477528919414874748dd363e6272efd5abf7654e68bef \
+    --hash=sha256:3031709084b6e7852d00479fd1d310b07d0ba82765f973b543c8af5061cf990e \
+    --hash=sha256:0689b5a8c5288bc0504d9fcee48f61a6a586b9b98514d7d29b840143d6734f39 \
+    --hash=sha256:7823bdd049099efa16e4246bdf15e5a13dbb18a51b68fa06d6c1d4d8b99a796e \
+    --hash=sha256:b792a349405fbc0163190fde0dc7b3fef3c9268292586cf5645598b48e63dc67 \
+    --hash=sha256:8373c6c251f7ef8bda6675dd6d2b3a0fcc31edf1201266b5cf608b62a37407f9
+
+usd-core==24.8 \
+    --hash=sha256:39fe8e266875e1105886cab870df4bcbe2d40a84696177bc574c22b96d843bf6 \
+    --hash=sha256:8b38b347dce9336d00dacd4c0e5813bcb43c61e165c653d754e3ee3dc4b2b715 \
+    --hash=sha256:d1dfe295ccbc57cac39e6dee2d7ce831d7a6a504b8a17dc974e631e3521f83a6 \
+    --hash=sha256:d3e06fd8c953c4de3d24591cdb9e8e65ca28a876519378d1e4a5d9bca419d7c1 \
+    --hash=sha256:7c18be89012d03f57d5445101695461c678f370d565259893c1cbcb71a6124e6 \
+    --hash=sha256:0a503029fa895c9d2307bde2b15748478a0bcb4e8c4ebfa90c7c613fb8137f9d \
+    --hash=sha256:a3e1d9b34d3936716dbce4ee22c135acffab0f58ba9da09d47dc74992cfc5259 \
+    --hash=sha256:c8efa7784fd11c3e97dec752d93913a234e976d1ebdf3ba8a23af606c86affb0 \
+    --hash=sha256:db4951f221323dea83795015f27105a375f28ddabe15ee237acc0728b8dd0212 \
+    --hash=sha256:0de3647478e7165d82d73e4b5eb30a7f83df167cc9b6cec3d19fd75fe6d01302 \
+    --hash=sha256:ffdba837680171a468d7ec29edc243fcd47d8d1e044fa810c4037b920d31ef60 \
+    --hash=sha256:20c45a3099662d32c1a059f3a08adc7829ef7ce40977c33e28835c6bb34aa321 \
+    --hash=sha256:6ecef8b1592a37c03c3cc3255f302cb2f5e67d4b9f691084b4674ed92032243c \
+    --hash=sha256:ba067c7476e509b49fa1cec3f19821ab8b10c9b43eeb6cae82676bb498297710 \
+    --hash=sha256:f3cafd0ddf647a1a8b2335ba109019628252d775c0b3d0f18407211b7801fe53

python/mujoco/testdata/usd_golden.usda

@@ -5,56 +5,74 @@
 """
     endTimeCode = 1
     startTimeCode = 0
-    timeCodesPerSecond = 24
+    timeCodesPerSecond = 60
     upAxis = "Z"
 )
 
 def Xform "World"
 {
     def Xform "Light_Xform_0"
     {
-        double3 xformOp:translate
+        double3 xformOp:translate.timeSamples = {
+            0: (0, -3.674234628677368, 3.674234628677368),
+        }
         uniform token[] xformOpOrder = ["xformOp:translate"]
 
         def SphereLight "Light_0"
         {
+            color3f inputs:color = (0.4, 0.4, 0.4)
+            float inputs:intensity = 10000
             bool inputs:normalize = 1
             float inputs:radius = 0.3
             bool treatAsPoint = 0
         }
     }
 
-    def Xform "CubeMesh_Xform_white_box"
+    def Xform "Camera_Xform_closeup"
     {
         matrix4d xformOp:transform.timeSamples = {
-            0: ( (0, 0, 0, 0), (0, 0, 0, 0), (0, 0, 0, 0), (0, 0, 0, 1) ),
+            0: ( (1, 0, -0, 0), (0, 0.009999499656260014, 0.9999499917030334, 0), (0, -0.9999499917030334, 0.009999499656260014, 0), (0, -6, 0, 1) ),
         }
         uniform token[] xformOpOrder = ["xformOp:transform"]
 
-        def Mesh "CubeMesh_white_box" (
+        def Camera "Camera_closeup"
+        {
+            float2 clippingRange = (0.0001, 1000000)
+            float focalLength = 12
+            float focusDistance = 400
+            float horizontalAperture = 12
+        }
+    }
+
+    def Xform "Mesh_Xform_white_box_id0_geom"
+    {
+        token visibility.timeSamples = {
+            0: "inherited",
+            1: "invisible",
+        }
+        float3 xformOp:scale
+        matrix4d xformOp:transform.timeSamples = {
+            0: ( (1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1) ),
+        }
+        uniform token[] xformOpOrder = ["xformOp:transform", "xformOp:scale"]
+
+        def Mesh "Mesh_white_box_id0_geom" (
             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>
+            rel material:binding = </World/_materials/Material_white_box_id0_geom>
             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"
+        def Material "Material_white_box_id0_geom"
         {
-            token outputs:surface.connect = </World/_materials/Material_white_box/Principled_BSDF.outputs:surface>
+            token outputs:surface.connect = </World/_materials/Material_white_box_id0_geom/Principled_BSDF.outputs:surface>
 
             def Shader "Principled_BSDF"
             {

python/mujoco/usd/demo.py

@@ -79,7 +79,10 @@ def generate_usd_trajectory(local_args):
   )
 
   parser.add_argument(
-      '--camera_names', type=str, nargs='+', help='cameras to include in usd'
+      '--camera_names',
+      type=str,
+      nargs='+',
+      help='cameras to include in usd'
   )
 
   parser.add_argument(

python/mujoco/usd/exporter.py

@@ -15,7 +15,7 @@
 """USD exporter."""
 
 import os
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import mujoco
 import mujoco.usd.camera as camera_module
@@ -25,8 +25,6 @@
 import numpy as np
 from PIL import Image as im
 from PIL import ImageOps
-import scipy
-import termcolor
 
 # TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
 # pytype: disable=module-attr
@@ -115,7 +113,7 @@ def __init__(
     self.usd_cameras = []
 
     # initializing rendering requirements
-    self.renderer = mujoco.Renderer(model, height, width, max_geom)
+    self._scene = mujoco.MjvScene(model=model, maxgeom=max_geom)
     self._initialize_usd_stage()
     self._scene_option = mujoco.MjvOption()  # using default scene option
 
@@ -133,7 +131,7 @@ def usd(self):
   @property
   def scene(self):
     """Returns the scene."""
-    return self.renderer.scene
+    return self._scene
 
   def _initialize_usd_stage(self):
     """Initializes a USD stage to represent the mujoco scene."""
@@ -165,14 +163,50 @@ def _initialize_output_directories(self):
       os.makedirs(self.assets_directory)
 
     if self.verbose:
-      print(
-          termcolor.colored(
-              "Writing output frames and assets to"
-              f" {self.output_directory_path}",
-              "green",
-          )
+      print(f"Writing output frames and assets to {self.output_directory_path}")
+
+  def _update_scene(
+      self,
+      data: mujoco.MjData,
+      camera: Union[int, str] = -1,
+      scene_option: Optional[mujoco.MjvOption] = None,
+  ):
+    """Updates the scene."""
+    camera_id = camera
+    if isinstance(camera_id, str):
+      camera_id = mujoco.mj_name2id(
+          self.model, mujoco.mjtObj.mjOBJ_CAMERA.value, camera_id
+      )
+      if camera_id == -1:
+        raise ValueError(f"The camera '{camera}' does not exist.")
+    if camera_id < -1 or camera_id >= self.model.ncam:
+      raise ValueError(
+          f"The camera id {camera_id} is out of range [-1, {self.model.ncam})."
       )
 
+    # Render camera.
+    camera = mujoco.MjvCamera()
+    camera.fixedcamid = camera_id
+
+    # Defaults to mjCAMERA_FREE, otherwise mjCAMERA_FIXED refers to a
+    # camera explicitly defined in the model.
+    if camera_id == -1:
+      camera.type = mujoco.mjtCamera.mjCAMERA_FREE
+      mujoco.mjv_defaultFreeCamera(self.model, camera)
+    else:
+      camera.type = mujoco.mjtCamera.mjCAMERA_FIXED
+
+    scene_option = scene_option or self._scene_option
+    mujoco.mjv_updateScene(
+        self.model,
+        data,
+        scene_option,
+        None,
+        camera,
+        mujoco.mjtCatBit.mjCAT_ALL.value,
+        self._scene,
+    )
+
   def update_scene(
       self,
       data: mujoco.MjData,
@@ -190,11 +224,10 @@ def update_scene(
     scene_option = scene_option or self._scene_option
 
     # update the mujoco renderer
-    self.renderer.update_scene(data, scene_option=scene_option)
+    self._update_scene(data, scene_option=scene_option)
 
     if self.updates == 0:
       self._initialize_usd_stage()
-
       self._load_lights()
       self._load_cameras()
 
@@ -237,11 +270,8 @@ def _load_textures(self):
 
     if self.verbose:
       print(
-          termcolor.colored(
-              f"Completed writing {self.model.ntex} textures to"
-              f" {self.assets_directory}",
-              "green",
-          )
+          f"Completed writing {self.model.ntex} textures to"
+          f" {self.assets_directory}",
       )
 
   def _load_geom(self, geom: mujoco.MjvGeom):
@@ -383,8 +413,10 @@ def _update_cameras(
       camera = self.usd_cameras[i]
       camera_name = self.camera_names[i]
 
-      self.renderer.update_scene(
-          data, scene_option=scene_option, camera=camera_name
+      self._update_scene(
+          data,
+          camera=camera_name,
+          scene_option=scene_option,
       )
 
       avg_camera = mujoco.mjv_averageCamera(
@@ -450,9 +482,11 @@ def add_camera(
     new_camera = camera_module.USDCamera(
         stage=self.stage, obj_name=obj_name)
 
-    r = scipy.spatial.transform.Rotation.from_euler(
-        "xyz", rotation_xyz, degrees=True)
-    new_camera.update(cam_pos=np.array(pos), cam_mat=r.as_matrix(), frame=0)
+    rotation = np.zeros(9)
+    quat = np.zeros(4)
+    mujoco.mju_euler2Quat(quat, rotation_xyz, "xyz")
+    mujoco.mju_quat2Mat(rotation, quat)
+    new_camera.update(cam_pos=np.array(pos), cam_mat=rotation, frame=0)
 
   def save_scene(self, filetype: str = "usd"):
     """Saves the scene to a USD file."""
@@ -468,11 +502,7 @@ def save_scene(self, filetype: str = "usd"):
         f"frames/frame_{self.frame_count}.{filetype}"
     )
     if self.verbose:
-      print(
-          termcolor.colored(
-              f"Completed writing frame_{self.frame_count}.{filetype}", "green"
-          )
-      )
+      print(f"Completed writing frame_{self.frame_count}.{filetype}")
 
   def _get_geom_name(self, geom):
     """Adding id as part of name for USD file."""

python/mujoco/usd/exporter_test.py

@@ -14,30 +14,23 @@
 # ==============================================================================
 """Tests for the MuJoCo USD Exporter."""
 
-import logging
 import os
+import tempfile
 
 from absl.testing import absltest
 from etils import epath
 import mujoco
-
-
-# Open3D and USD are not fully supported on all MuJoCo architectures.
-execute_test = True
-try:
-  from mujoco.usd import exporter as exporter_module  # pylint: disable=g-import-not-at-top
-except ImportError:
-  logging.warning('Skipping test due to missing import')
-  execute_test = False
+from mujoco.usd import exporter as exporter_module  # pylint: disable=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')
+    output_dir_root = os.getenv(
+        "TEST_UNDECLARED_OUTPUTS_DIR", tempfile.gettempdir()
+    )
+    output_dir_name = "usd_test"
     xml = """
 <mujoco>
   <worldbody>
@@ -50,19 +43,27 @@ def test_usd_export(self):
     data = mujoco.MjData(model)
     exporter = exporter_module.USDExporter(
         model,
-        output_directory_name='mujoco_usdpkg',
-        output_directory_root=output_dir,
+        output_directory_name=output_dir_name,
+        output_directory_root=output_dir_root,
+        camera_names=["closeup"],
     )
+    mujoco.mj_step(model, data)
     exporter.update_scene(data)
-    exporter.save_scene('export.usda')
+    exporter.save_scene("usda")
 
-    with open(os.path.join(
-        output_dir, 'mujoco_usdpkg/frames', 'frame_1.export.usda'), 'r') as f:
+    with open(
+        os.path.join(
+            output_dir_root, f"{output_dir_name}/frames", "frame_1.usda"
+        ),
+        "r",
+        encoding="utf-8",
+    ) 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())
+          epath.resource_path("mujoco"), "testdata", "usd_golden.usda"
+      )
+      with open(golden_path, "r", encoding="utf-8") as golden_file:
+        self.assertEqual(f.readlines(), golden_file.readlines())
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   absltest.main()