From 7a1e2bf61e60ef38061b9ca352213a91123543b4 Mon Sep 17 00:00:00 2001
From: Abhishek Joshi <abhishekjoshinj@gmail.com>
Date: Thu, 2 May 2024 17:23:48 -0500
Subject: [PATCH] adding demo file and unnamed geom fix

---
 python/mujoco/usd/component.py |  2 +-
 python/mujoco/usd/demo.py      | 53 ++++++++++++++++++++++++++++++----
 python/mujoco/usd/exporter.py  | 25 +++++++++++-----
 3 files changed, 67 insertions(+), 13 deletions(-)

diff --git a/python/mujoco/usd/component.py b/python/mujoco/usd/component.py
index 48b3e251..e91fc8e2 100644
--- a/python/mujoco/usd/component.py
+++ b/python/mujoco/usd/component.py
@@ -16,7 +16,7 @@ from typing import List, Optional, Tuple
 
 import mujoco
 
-import mujoco.usd.utils
+import mujoco.usd.utils as utils_component
 import mujoco.usd.shapes as shapes_component
 
 import numpy as np
diff --git a/python/mujoco/usd/demo.py b/python/mujoco/usd/demo.py
index e8e8deb1..13b6f21f 100644
--- a/python/mujoco/usd/demo.py
+++ b/python/mujoco/usd/demo.py
@@ -1,20 +1,63 @@
+import argparse
+from tqdm import tqdm
+from pathlib import Path
 
 import mujoco
 from mujoco.usd import exporter
 
-if __name__ == "__main__":
+def generate_usd_trajectory(args):
 
   # load a model to mujoco
-  m = mujoco.MjModel.from_xml_path("/Users/abhishek/Documents/research/mujoco/test/engine/testdata/catenary.xml")
+  model_path = args.model_path
+  m = mujoco.MjModel.from_xml_path(model_path)
   d = mujoco.MjData(m)
 
   # create an instance of the USDExporter
-  exp = exporter.USDExporter(model=m)
+  exp = exporter.USDExporter(model=m,
+                             output_directory_name=Path(args.model_path).stem,
+                             output_directory_root=args.output_directory_root,
+                             camera_names=args.camera_names)
 
-  for i in range(100):
+  # step through the model for length steps
+  for i in tqdm(range(args.length)):
     mujoco.mj_step(m, d)
     exp.update_scene(d)
 
-  exp.save_scene(filetype="usd")
+  exp.save_scene(filetype=args.export_extension)
+
+
+if __name__ == "__main__":
+
+  parser = argparse.ArgumentParser()
+
+  parser.add_argument('--model_path', 
+                      type=str,
+                      required=True,
+                      help='path to mjcf xml model')
+
+  parser.add_argument('--length',
+                      type=int,
+                      default=100,
+                      help='length of trajectory to render')
+
+  parser.add_argument('--output_directory_root', 
+                      type=str,
+                      default="../usd_trajectories/",
+                      help='location where to create usd files')
+
+  parser.add_argument('--camera_names', 
+                      type=str,
+                      nargs='+',
+                      help='cameras to include in usd')
+
+  parser.add_argument('--export_extension', 
+                      type=str,
+                      default="usd",
+                      help='extension of exported file (can be usd, usda, or usdc)')
+
+  args = parser.parse_args()
+  generate_usd_trajectory(args)
+
+  
 
 
diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py
index ff5b24f3..d4b03fae 100644
--- a/python/mujoco/usd/exporter.py
+++ b/python/mujoco/usd/exporter.py
@@ -128,7 +128,7 @@ class USDExporter:
     UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
     self.stage.SetStartTimeCode(0)
     # add as user input
-    self.stage.SetTimeCodesPerSecond(24.0)
+    self.stage.SetTimeCodesPerSecond(60.0)
 
     default_prim = UsdGeom.Xform.Define(
         self.stage, Sdf.Path("/World")
@@ -232,9 +232,7 @@ class USDExporter:
 
   def _load_geom(self, geom: mujoco.MjvGeom):
 
-    geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
-    if not geom_name:
-      geom_name = f"unnamed_geom{geom.objid}"
+    geom_name = self._get_geom_name(geom.objtype, geom.objid)
 
     assert geom_name not in self.geom_name2usd
 
@@ -271,12 +269,17 @@ class USDExporter:
 
     geom_names = set(self.geom_name2usd.keys())
 
-    # iterate through all geoms in the scene and makes update
     for i in range(self.scene.ngeom):
       geom = self.scene.geoms[i]
       geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
       if not geom_name:
-        geom_name = f"unnamed_geom{geom.objid}"
+        geom_name = "None"
+      geom_name += f"_{geom.objid}"
+
+    # iterate through all geoms in the scene and makes update
+    for i in range(self.scene.ngeom):
+      geom = self.scene.geoms[i]
+      geom_name = self._get_geom_name(geom.objtype, geom.objid)
 
       if geom_name not in self.geom_name2usd:
         self._load_geom(geom)
@@ -394,9 +397,17 @@ class USDExporter:
     new_camera.update(cam_pos=np.array(pos), cam_mat=r.as_matrix(), frame=0)
 
   def save_scene(self, filetype: str = "usd"):
+    assert filetype in ["usd", "usda", "usdc"]
     self.stage.SetEndTimeCode(self.frame_count)
     self.stage.Export(
         f"{self.output_directory_root}/{self.output_directory_name}/frames/frame_{self.frame_count}_.{filetype}"
     )
     if self.verbose:
-      print(termcolor.colored(f"Writing frame_{self.frame_count}", "green"))
+      print(termcolor.colored(f"Completed writing frame_{self.frame_count}.{filetype}", "green"))
+
+  def _get_geom_name(self, objtype, objid):
+    geom_name = mujoco.mj_id2name(self.model, objtype, objid)
+    if not geom_name:
+      geom_name = "None"
+    geom_name += f"_{objid}"
+    return geom_name