fixing frame rate bug in demo file

python/mujoco/usd/camera.py

@@ -18,7 +18,8 @@ from pxr import Gf
 from pxr import Usd
 from pxr import UsdGeom
 
-import mujoco.usd.utils as utils_component
+# import mujoco.usd.utils as utils_component
+import utils as utils_component
 
 class USDCamera:
   """Class that handles the cameras in the USD scene"""

python/mujoco/usd/demo.py

@@ -1,26 +1,14 @@
-# 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.
-# ==============================================================================
 import argparse
-from pathlib import Path
 from tqdm import tqdm
+from pathlib import Path
 
 import mujoco
-from mujoco.usd import exporter
+
+# from mujoco.usd import exporter
+import exporter
 
 def generate_usd_trajectory(args):
-  """Generates a USD file from a mujoco trajectory."""
+
   # load a model to mujoco
   model_path = args.model_path
   m = mujoco.MjModel.from_xml_path(model_path)
@@ -31,12 +19,12 @@ def generate_usd_trajectory(args):
                              output_directory_name=Path(args.model_path).stem,
                              output_directory_root=args.output_directory_root,
                              camera_names=args.camera_names)
-
-  # step through the model for length steps
-  for _ in tqdm(range(args.length)):
-    for _ in range(args.steps_per_frame):
-      mujoco.mj_step(m, d)
-    exp.update_scene(d)
+    
+  # step through the simulation for the given duration of time
+  while d.time < args.duration:
+    mujoco.mj_step(m, d)
+    if exp.frame_count < d.time * args.framerate:
+      exp.update_scene(data=d)
 
   exp.save_scene(filetype=args.export_extension)
 
@@ -44,35 +32,35 @@ if __name__ == "__main__":
 
   parser = argparse.ArgumentParser()
 
-  parser.add_argument('--model_path',
+  parser.add_argument('--model_path', 
                       type=str,
                       required=True,
                       help='path to mjcf xml model')
 
-  parser.add_argument('--length',
+  parser.add_argument('--duration',
                       type=int,
-                      default=100,
-                      help='length of trajectory to render')
+                      default=5,
+                      help='duration in seconds for the generated video')
 
-  parser.add_argument('--output_directory_root',
+  parser.add_argument('--framerate',
+                      type=int,
+                      default=60,
+                      help='frame rate of the generated video')
+
+  parser.add_argument('--output_directory_root', 
                       type=str,
                       default="../usd_trajectories/",
                       help='location where to create usd files')
 
-  parser.add_argument('--camera_names',
+  parser.add_argument('--camera_names', 
                       type=str,
                       nargs='+',
                       help='cameras to include in usd')
 
-  parser.add_argument('--export_extension',
+  parser.add_argument('--export_extension', 
                       type=str,
                       default="usd",
-                      help='extension of exported file (usd, usda, or usdc)')
-
-  parser.add_argument('--steps_per_frame',
-                      type=int,
-                      default=1,
-                      help='number of frames to skip for each rendering step')
+                      help='extension of exported file (can be usd, usda, or usdc)')
 
   args = parser.parse_args()
   generate_usd_trajectory(args)

python/mujoco/usd/exporter.py

@@ -29,10 +29,14 @@ from pxr import UsdGeom
 
 import mujoco
 
-import mujoco.usd.shapes as shapes_module
-import mujoco.usd.objects as object_module
-import mujoco.usd.lights as light_module
-import mujoco.usd.camera as camera_module
+# import mujoco.usd.shapes as shapes_module
+# import mujoco.usd.objects as object_module
+# import mujoco.usd.lights as light_module
+# import mujoco.usd.camera as camera_module
+import shapes as shapes_module
+import objects as object_module
+import lights as light_module
+import camera as camera_module
 
 
 class USDExporter:
@@ -293,7 +297,6 @@ class USDExporter:
     self.geom_refs[geom_name] = usd_geom
 
   def _update_geoms(self):
-
     # iterate through all geoms in the scene and makes update
     for i in range(self.scene.ngeom):
       geom = self.scene.geoms[i]

python/mujoco/usd/exporter_test.py

@@ -59,7 +59,7 @@ class ExporterTest(absltest.TestCase):
     exporter.save_scene("export.usda")
 
     with open(os.path.join(
-        output_dir, "mujoco_usdpkg/frames", "frame_1_.export.usda"), "r") as f:
+        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:

python/mujoco/usd/objects.py

@@ -27,8 +27,11 @@ from pxr import Vt
 
 import mujoco
 
-import mujoco.usd.utils as utils_component
-import mujoco.usd.shapes as shapes_component
+# import mujoco.usd.utils as utils_component
+# import mujoco.usd.shapes as shapes_component
+import utils as utils_component
+import shapes as shapes_component
+
 
 class USDObject(ABC):
   """ Abstract interface for all USD objects including meshes and primitives.

python/mujoco/usd/shapes.py

@@ -151,7 +151,7 @@ def mesh_generator(
     prim_mesh, mesh = None, None
 
     for shape, config in mesh_config.items():
-
+      
       if "name" in shape:
         continue