updating geom visibility logic; fixing infinite plane texture bug; adding frame skip arg to demo

python/mujoco/usd/component.py

@@ -16,8 +16,10 @@ from typing import List, Optional, Tuple
 
 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
 
 import numpy as np
 
@@ -80,6 +82,8 @@ class USDMesh:
     # defining ops required by update function
     self.transform_op = self.usd_xform.AddTransformOp()
 
+    self.last_visible_frame = -2 # not an arbitary value, forces difference greater than 1 for visibility on 0th frame
+
   def get_facetexcoord_ranges(self, nmesh, arr):
     facetexcoords_ranges = [0]
     running_sum = 0
@@ -224,13 +228,18 @@ class USDMesh:
         rotation_matrix=mat, translation_vector=pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
-    self.update_visibility(visible, frame)
+    
+    if visible and frame - self.last_visible_frame > 1:
+      # non consecutive visible frames
+      self.update_visibility(False, max(0, self.last_visible_frame))
+      self.update_visibility(True, frame)
+
+    if visible:
+      self.last_visible_frame = frame
 
   def update_visibility(self, visible: bool, frame: int):
-    if visible:
-      self.usd_prim.GetAttribute("visibility").Set("inherited", frame)
-    else:
-      self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
+    visibility_setting = "inherited" if visible else "invisible"
+    self.usd_prim.GetAttribute("visibility").Set(visibility_setting, frame)
 
 
 class USDPrimitiveMesh:
@@ -290,6 +299,8 @@ class USDPrimitiveMesh:
     # defining ops required by update function
     self.transform_op = self.usd_xform.AddTransformOp()
 
+    self.last_visible_frame = -2
+
   def _set_refinement_properties(self):
     self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
 
@@ -298,14 +309,17 @@ class USDPrimitiveMesh:
     assert self.prim_mesh
 
     x_scale, y_scale = self.geom.texrepeat
-
     mesh_texcoord = np.array(self.prim_mesh.triangle_uvs)
     mesh_facetexcoord = np.asarray(self.prim_mesh.triangles)
 
     x_multiplier, y_multiplier = 1, 1
     if self.geom.texuniform:
-      x_multiplier, y_multiplier = self.geom.size[:2]
-
+      # x_multiplier, y_multiplier = self.geom.size[:2]
+      if "box" in self.mesh_config:
+        x_multiplier, y_multiplier = self.mesh_config["box"]["width"], self.mesh_config["box"]["height"]
+      elif "sphere" in self.mesh_config:
+        x_multiplier, y_multiplier = self.mesh_config["sphere"]["radius"], self.mesh_config["sphere"]["radius"]
+      
     mesh_texcoord[:, 0] *= x_scale * x_multiplier
     mesh_texcoord[:, 1] *= y_scale * y_multiplier
 
@@ -408,13 +422,18 @@ class USDPrimitiveMesh:
         rotation_matrix=mat, translation_vector=pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
-    self.update_visibility(visible, frame)
+
+    if visible and frame - self.last_visible_frame > 1:
+      # non consecutive visible frames
+      self.update_visibility(False, max(0, self.last_visible_frame))
+      self.update_visibility(True, frame)
+
+    if visible:
+      self.last_visible_frame = frame
 
   def update_visibility(self, visible: bool, frame: int):
-    if visible:
-      self.usd_prim.GetAttribute("visibility").Set("inherited", frame)
-    else:
-      self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
+    visibility_setting = "inherited" if visible else "invisible"
+    self.usd_prim.GetAttribute("visibility").Set(visibility_setting, frame)
 
 class USDSphereLight:
 

python/mujoco/usd/demo.py

@@ -3,7 +3,8 @@ from tqdm import tqdm
 from pathlib import Path
 
 import mujoco
-from mujoco.usd import exporter
+# from mujoco.usd import exporter
+import exporter as exporter
 
 def generate_usd_trajectory(args):
 
@@ -20,12 +21,12 @@ def generate_usd_trajectory(args):
 
   # step through the model for length steps
   for i in tqdm(range(args.length)):
-    mujoco.mj_step(m, d)
+    for i in range(args.steps_per_frame):
+      mujoco.mj_step(m, d)
     exp.update_scene(d)
 
   exp.save_scene(filetype=args.export_extension)
 
-
 if __name__ == "__main__":
 
   parser = argparse.ArgumentParser()
@@ -55,6 +56,11 @@ if __name__ == "__main__":
                       default="usd",
                       help='extension of exported file (can be usd, usda, or usdc)')
 
+  parser.add_argument('--steps_per_frame',
+                      type=int,
+                      default=1,
+                      help='number of frames to skip for each rendering step')
+
   args = parser.parse_args()
   generate_usd_trajectory(args)
 

python/mujoco/usd/exporter.py

@@ -16,8 +16,10 @@ import os
 
 import mujoco
 
-import mujoco.usd.shapes as shapes_module
-import mujoco.usd.component as component_module
+# import mujoco.usd.shapes as shapes_module
+# import mujoco.usd.component as component_module
+import shapes as shapes_module
+import component as component_module
 
 import numpy as np
 import scipy
@@ -102,7 +104,8 @@ class USDExporter:
     self.frame_count = 0  # maintains how many times we have saved the scene
     self.updates = 0
 
-    self.geom_name2usd = {}
+    self.geom_names = set()
+    self.geom_refs = {}
 
     # initializing rendering requirements
     self.renderer = mujoco.Renderer(model, height, width, max_geom)
@@ -215,7 +218,7 @@ class USDExporter:
       )
       img_path = os.path.join(
           relative_path, texture_file_name
-      )  # relative path, TODO: switch back to this
+      )
 
       self.texture_files.append(img_path)
 
@@ -232,9 +235,9 @@ class USDExporter:
 
   def _load_geom(self, geom: mujoco.MjvGeom):
 
-    geom_name = self._get_geom_name(geom.objtype, geom.objid)
+    geom_name = self._get_geom_name(geom)
 
-    assert geom_name not in self.geom_name2usd
+    assert geom_name not in self.geom_names
 
     texture_file = self.texture_files[geom.texid] if geom.texid != -1 else None
 
@@ -263,42 +266,30 @@ class USDExporter:
           texture_file=texture_file,
       )
 
-    self.geom_name2usd[geom_name] = usd_geom
+    self.geom_names.add(geom_name)
+    self.geom_refs[geom_name] = usd_geom
 
   def _update_geoms(self):
 
-    geom_names = set(self.geom_name2usd.keys())
-
-    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 = "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)
+      geom_name = self._get_geom_name(geom)
 
-      if geom_name not in self.geom_name2usd:
-        self._load_geom(geom)
-        if self.geom_name2usd[geom_name]:
-          self.geom_name2usd[geom_name].update_visibility(False, 0)
+      if "tendon" in geom_name:
+        # we handle tendon geoms separately to avoid naming conflicts
+        pass
+      else:
+        if geom_name not in self.geom_names:
+          # load a new object into USD
+          self._load_geom(geom)
 
-      if self.geom_name2usd[geom_name]:
-        self.geom_name2usd[geom_name].update(
+        self.geom_refs[geom_name].update(
             pos=geom.pos,
             mat=geom.mat,
             visible=geom.rgba[3] > 0,
             frame=self.updates,
         )
-      if geom_name in geom_names:
-        geom_names.remove(geom_name)
-
-    for geom_name in geom_names:
-      if self.geom_name2usd[geom_name]:
-        self.geom_name2usd[geom_name].update_visibility(False, self.updates)
 
   def _load_lights(self):
     # initializes an usd light object for every light in the scene
@@ -399,15 +390,35 @@ class USDExporter:
   def save_scene(self, filetype: str = "usd"):
     assert filetype in ["usd", "usda", "usdc"]
     self.stage.SetEndTimeCode(self.frame_count)
+
+    # post-processing for visibility of geoms in scene
+    for _, geom_ref in self.geom_refs.items():
+      geom_ref.update_visibility(False, geom_ref.last_visible_frame+1)
+
     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"Completed writing frame_{self.frame_count}.{filetype}", "green"))
 
-  def _get_geom_name(self, objtype, objid):
-    geom_name = mujoco.mj_id2name(self.model, objtype, objid)
+  def _get_geom_name(self, geom):
+    # adding id as part of name for USD file
+    geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
     if not geom_name:
       geom_name = "None"
-    geom_name += f"_{objid}"
+    geom_name += f"_id{geom.objid}"
+
+    # adding additional naming information to differentiate between geoms and tendons
+    if geom.objtype == mujoco.mjtObj.mjOBJ_GEOM:
+      geom_name += "_geom"
+    elif geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
+      geom_name += "_tendon"
+
     return geom_name
+
+  # for debugging purposes, prints all geoms in scene including those part of tendons
+  def _print_scene_geom_info(self):
+    for i in range(self.scene.ngeom):
+      geom = self.scene.geoms[i]
+      geom_name = self._get_geom_name(geom)
+      print(i, geom_name)