From bb2c7ae556cc8b3a0c8fdba42ad3b51d4ee67235 Mon Sep 17 00:00:00 2001
From: Abhishek Joshi <abhishekjoshinj@gmail.com>
Date: Thu, 9 May 2024 08:42:29 -0500
Subject: [PATCH] adding initial tendon logic

---
 python/mujoco/usd/component.py | 35 ++++++++++++++---
 python/mujoco/usd/exporter.py  | 68 +++++++++++++++++++++++-----------
 2 files changed, 76 insertions(+), 27 deletions(-)

diff --git a/python/mujoco/usd/component.py b/python/mujoco/usd/component.py
index d7729e52..28fa3efd 100644
--- a/python/mujoco/usd/component.py
+++ b/python/mujoco/usd/component.py
@@ -296,8 +296,9 @@ class USDPrimitiveMesh:
     # setting attributes for the shape
     self._attach_material()
 
-    # defining ops required by update function
+    # defining ops required by update functions
     self.transform_op = self.usd_xform.AddTransformOp()
+    self.scale_op = self.usd_xform.AddScaleOp()
 
     self.last_visible_frame = -2
 
@@ -314,7 +315,6 @@ class USDPrimitiveMesh:
 
     x_multiplier, y_multiplier = 1, 1
     if self.geom.texuniform:
-      # 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:
@@ -435,6 +435,33 @@ class USDPrimitiveMesh:
     visibility_setting = "inherited" if visible else "invisible"
     self.usd_prim.GetAttribute("visibility").Set(visibility_setting, frame)
 
+  def update_scale(self, scale: np.ndarray, frame: int):
+    self.scale_op.Set(Gf.Vec3f(scale.tolist()), frame)
+
+class USDTendon(USDPrimitiveMesh):
+
+  def __init__(
+    self,
+    mesh_config: List[dict],
+    stage: Usd.Stage,
+    geom: mujoco.MjvGeom,
+    obj_name: str,
+    rgba: np.ndarray = np.array([1, 1, 1, 1]),
+    texture_file: Optional[str] = None,
+  ):
+    super().__init__(
+      mesh_config=mesh_config,
+      stage=stage,
+      geom=geom,
+      obj_name=obj_name,
+      rgba=rgba,
+      texture_file=texture_file
+    )
+
+  def update(self, pos: np.ndarray, mat: np.ndarray, scale: np.ndarray, visible: bool, frame: int):
+    super().update(pos, mat, visible, frame)
+    super().update_scale(scale, frame)
+
 class USDSphereLight:
 
   def __init__(
@@ -456,9 +483,7 @@ class USDSphereLight:
     # defining ops required by update function
     self.translate_op = self.usd_xform.AddTranslateOp()
 
-  def update(
-      self, pos: np.ndarray, intensity: int, color: np.ndarray, frame: int
-  ):
+  def update(self, pos: np.ndarray, intensity: int, color: np.ndarray, frame: int):
     self.translate_op.Set(Gf.Vec3d(pos.tolist()), frame)
 
     if not np.any(pos):
diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py
index 745f430a..69129e63 100644
--- a/python/mujoco/usd/exporter.py
+++ b/python/mujoco/usd/exporter.py
@@ -233,7 +233,7 @@ class USDExporter:
           )
       )
 
-  def _load_geom(self, geom: mujoco.MjvGeom):
+  def _load_geom(self, geom: mujoco.MjvGeom, tendon: Optional[bool]=False):
 
     geom_name = self._get_geom_name(geom)
 
@@ -252,19 +252,35 @@ class USDExporter:
           texture_file=texture_file,
       )
     else:
-      mesh_config = shapes_module.mesh_config_generator(
-          name=geom_name,
-          geom_type=geom.type,
-          size=geom.size   
-      )
-      usd_geom = component_module.USDPrimitiveMesh(
-          mesh_config=mesh_config,
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
+      if geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
+        # scale of a tendon is equivalent to its size, assume a unit length geom size
+        mesh_config = shapes_module.mesh_config_generator(
+            name=geom_name,
+            geom_type=geom.type,
+            size=np.array([1.0, 1.0, 1.0])
+        )
+        usd_geom = component_module.USDTendon(
+            mesh_config=mesh_config,
+            stage=self.stage,
+            geom=geom,
+            obj_name=geom_name,
+            rgba=geom.rgba,
+            texture_file=texture_file,
+        )
+      else:
+        mesh_config = shapes_module.mesh_config_generator(
+            name=geom_name,
+            geom_type=geom.type,
+            size=geom.size   
+        )
+        usd_geom = component_module.USDPrimitiveMesh(
+            mesh_config=mesh_config,
+            stage=self.stage,
+            geom=geom,
+            obj_name=geom_name,
+            rgba=geom.rgba,
+            texture_file=texture_file,
+        )
 
     self.geom_names.add(geom_name)
     self.geom_refs[geom_name] = usd_geom
@@ -276,14 +292,22 @@ class USDExporter:
       geom = self.scene.geoms[i]
       geom_name = self._get_geom_name(geom)
 
-      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 geom_name not in self.geom_names:
+        # load a new object into USD
+        self._load_geom(geom)
 
+      # TODO: make all objects scale values equivalent to their size, set original object mehs to have unit dimensions
+      if geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
+        tendon_scale = geom.size
+        tendon_scale[2] /= 2
+        self.geom_refs[geom_name].update(
+            pos=geom.pos,
+            mat=geom.mat,
+            scale=tendon_scale,
+            visible=geom.rgba[3] > 0,
+            frame=self.updates,
+        )
+      else:
         self.geom_refs[geom_name].update(
             pos=geom.pos,
             mat=geom.mat,
@@ -412,7 +436,7 @@ class USDExporter:
     if geom.objtype == mujoco.mjtObj.mjOBJ_GEOM:
       geom_name += "_geom"
     elif geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
-      geom_name += "_tendon"
+      geom_name += f"_tendon_segid{geom.segid}"
 
     return geom_name