decoupling tendons to subcomponents and fixing tendon scaling issue

python/mujoco/usd/__init__.py

@@ -1,4 +1,6 @@
 from .exporter import *
-from .component import *
+from .objects import *
+from .lights import *
+from .camera import *
 from .utils import *
 from .shapes import *

python/mujoco/usd/camera.py

@@ -0,0 +1,38 @@
+from typing import List, Optional, Tuple
+
+import utils as utils_component
+
+import numpy as np
+
+from pxr import Gf
+from pxr import Usd
+from pxr import UsdGeom
+
+class USDCamera:
+
+  def __init__(self, stage: Usd.Stage, obj_name: str):
+    self.stage = stage
+
+    xform_path = f"/World/Camera_Xform_{obj_name}"
+    camera_path = f"{xform_path}/Camera_{obj_name}"
+    self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
+    self.usd_camera = UsdGeom.Camera.Define(stage, camera_path)
+    self.usd_prim = stage.GetPrimAtPath(camera_path)
+
+    # defining ops required by update function
+    self.transform_op = self.usd_xform.AddTransformOp()
+
+    # self.usd_camera.CreateFocalLengthAttr().Set(18.14756) # default in omniverse
+    self.usd_camera.CreateFocalLengthAttr().Set(12)
+    self.usd_camera.CreateFocusDistanceAttr().Set(400)
+
+    self.usd_camera.GetHorizontalApertureAttr().Set(12)
+
+    self.usd_camera.GetClippingRangeAttr().Set(Gf.Vec2f(1e-4, 1e6))
+
+  def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
+
+    transformation_mat = utils_component.create_transform_matrix(
+        rotation_matrix=cam_mat, translation_vector=cam_pos
+    ).T
+    self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)

python/mujoco/usd/exporter.py

@@ -17,9 +17,11 @@ import os
 import mujoco
 
 # import mujoco.usd.shapes as shapes_module
-# import mujoco.usd.component as component_module
+# import mujoco.usd.component as object_module
 import shapes as shapes_module
-import component as component_module
+import objects as object_module
+import lights as light_module
+import camera as camera_module
 
 import numpy as np
 import scipy
@@ -181,7 +183,6 @@ class USDExporter:
     # update the mujoco renderer
     self.renderer.update_scene(data, scene_option=scene_option)
 
-    # TODO: update scene options
     if self.updates == 0:
       self._initialize_usd_stage()
 
@@ -195,7 +196,6 @@ class USDExporter:
     self.updates += 1
 
   def _load_textures(self):
-    # TODO: remove code once added internally to mujoco
     data_adr = 0
     self.texture_files = []
     for texture_id in tqdm.tqdm(range(self.model.ntex)):
@@ -241,8 +241,9 @@ class USDExporter:
 
     texture_file = self.texture_files[geom.texid] if geom.texid != -1 else None
 
+    # handling meshes in our scene
     if geom.type == mujoco.mjtGeom.mjGEOM_MESH:
-      usd_geom = component_module.USDMesh(
+      usd_geom = object_module.USDMesh(
           stage=self.stage,
           model=self.model,
           geom=geom,
@@ -252,14 +253,15 @@ class USDExporter:
           texture_file=texture_file,
       )
     else:
+      # handling tendons in our scene
       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])
+            size=np.array([1.0, 1.0, 1.0]),
+            decouple=True
         )
-        usd_geom = component_module.USDTendon(
+        usd_geom = object_module.USDTendon(
             mesh_config=mesh_config,
             stage=self.stage,
             geom=geom,
@@ -267,13 +269,14 @@ class USDExporter:
             rgba=geom.rgba,
             texture_file=texture_file,
         )
+      # handling primitives in our scene
       else:
         mesh_config = shapes_module.mesh_config_generator(
             name=geom_name,
             geom_type=geom.type,
             size=geom.size   
         )
-        usd_geom = component_module.USDPrimitiveMesh(
+        usd_geom = object_module.USDPrimitiveMesh(
             mesh_config=mesh_config,
             stage=self.stage,
             geom=geom,
@@ -298,7 +301,6 @@ class USDExporter:
 
       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,
@@ -321,7 +323,7 @@ class USDExporter:
       light = self.scene.lights[i]
       if not np.allclose(light.pos, [0, 0, 0]):
         self.usd_lights.append
-        (component_module.USDSphereLight(stage=self.stage, obj_name=str(i)))
+        (light_module.USDSphereLight(stage=self.stage, obj_name=str(i)))
       else:
         self.usd_lights.append(None)
 
@@ -347,7 +349,7 @@ class USDExporter:
     if self.camera_names is not None:
       for name in self.camera_names:
         self.usd_cameras.append(
-            component_module.USDCamera(stage=self.stage, obj_name=name))
+            camera_module.USDCamera(stage=self.stage, obj_name=name))
 
   def _update_cameras(
       self,
@@ -388,11 +390,11 @@ class USDExporter:
   ):
 
     if light_type == "sphere":
-      new_light = component_module.USDSphereLight(stage=self.stage, obj_name=str(objid), radius=radius)
+      new_light = light_module.USDSphereLight(stage=self.stage, obj_name=str(objid), radius=radius)
 
       new_light.update(pos=np.array(pos), intensity=intensity, color=color, frame=0)
     elif light_type == "dome":
-      new_light = component_module.USDDomeLight(
+      new_light = light_module.USDDomeLight(
           stage=self.stage, obj_name=str(objid))
 
       new_light.update(intensity=intensity, color=color, frame=0)
@@ -403,7 +405,7 @@ class USDExporter:
       rotation_xyz: List[float],
       obj_name: Optional[str] = "camera_1",
   ):
-    new_camera = component_module.USDCamera(
+    new_camera = camera_module.USDCamera(
         stage=self.stage, obj_name=str(objid))
 
     r = scipy.spatial.transform.Rotation.from_euler(

python/mujoco/usd/lights.py

@@ -0,0 +1,59 @@
+from typing import List, Optional, Tuple
+
+import numpy as np
+
+from pxr import Gf
+from pxr import Usd
+from pxr import UsdGeom
+from pxr import UsdLux
+
+class USDSphereLight:
+
+  def __init__(
+      self, stage: Usd.Stage, obj_name: str, radius: Optional[float] = 0.3
+  ):
+    self.stage = stage
+
+    xform_path = f"/World/Light_Xform_{obj_name}"
+    light_path = f"{xform_path}/Light_{obj_name}"
+    self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
+    self.usd_light = UsdLux.SphereLight.Define(stage, light_path)
+    self.usd_prim = stage.GetPrimAtPath(light_path)
+
+    # we assume in mujoco that all lights are point lights
+    self.usd_light.GetRadiusAttr().Set(radius)
+    self.usd_light.GetTreatAsPointAttr().Set(False)
+    self.usd_light.GetNormalizeAttr().Set(True)
+
+    # 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):
+    self.translate_op.Set(Gf.Vec3d(pos.tolist()), frame)
+
+    if not np.any(pos):
+      intensity = 0
+
+    self.usd_light.GetIntensityAttr().Set(intensity)
+    self.usd_light.GetColorAttr().Set(Gf.Vec3d(color.tolist()))
+
+
+class USDDomeLight:
+
+  def __init__(self, stage: Usd.Stage, obj_name: str):
+    self.stage = stage
+
+    xform_path = f"/World/Light_Xform_{obj_name}"
+    light_path = f"{xform_path}/Light_{obj_name}"
+    self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
+    self.usd_light = UsdLux.DomeLight.Define(stage, light_path)
+    self.usd_prim = stage.GetPrimAtPath(light_path)
+
+    # we assume in mujoco that all lights are point lights
+    self.usd_light.GetNormalizeAttr().Set(True)
+
+  def update(self, intensity: int, color: np.ndarray, frame: int):
+    self.usd_light.GetIntensityAttr().Set(intensity)
+    self.usd_light.GetExposureAttr().Set(0.0)
+    self.usd_light.GetColorAttr().Set(Gf.Vec3d(color.tolist()))
+

python/mujoco/usd/objects.py

@@ -0,0 +1,497 @@
+# 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.
+# ==============================================================================
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import List, Optional, Tuple
+import pprint
+
+import mujoco
+
+# 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
+
+# TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
+# pytype: disable=module-attr
+from open3d import open3d as o3d
+from pxr import Gf
+from pxr import Sdf
+from pxr import Usd
+from pxr import UsdGeom
+from pxr import UsdLux
+from pxr import UsdShade
+from pxr import Vt
+
+class USDObject(ABC):
+  """ Abstract interface for all USD objects including meshes and primitives
+
+  Subclasses must implement:
+
+  * `_get_uv_geometry(self)`: gets the nessecary UV information to wrap a texture 
+     around an object in USD. Each subclass implements their own method to getting
+     UV information as different objects are contructed in different ways.
+
+  * `_get_mesh_geometry(self)`: gets the mesh geometry of an object in the scene.
+  """
+
+  def __init__(
+      self,
+      stage: Usd.Stage,
+      geom: mujoco.MjvGeom,
+      obj_name: str,
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
+      texture_file: Optional[str] = None,
+  ):
+    self.stage = stage
+    self.geom = geom
+    self.obj_name = obj_name
+    self.rgba = rgba
+    self.texture_file = texture_file
+
+    self.xform_path = f"/World/Mesh_Xform_{obj_name}"
+    self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path)
+
+    # defining ops required by update function
+    self.transform_op = self.usd_xform.AddTransformOp()
+    self.scale_op = self.usd_xform.AddScaleOp()
+
+    self.last_visible_frame = -2 # not an arbitary value, forces difference greater than 1 for visibility on 0th frame
+
+  @abstractmethod
+  def _get_uv_geometry(self):
+    """Gets UV information for an object in the scene"""
+    raise NotImplementedError
+  
+  @abstractmethod
+  def _get_mesh_geometry(self):
+    """Gets structure of an object in the scene"""
+    raise NotImplementedError
+
+  def attach_image_material(self, usd_mesh):
+    """Attaches an image texture to a material for a USD object"""
+    mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}")
+    mtl = UsdShade.Material.Define(self.stage, mtl_path)
+
+    bsdf_shader = UsdShade.Shader.Define(
+        self.stage, mtl_path.AppendPath("Principled_BSDF")
+    )
+    image_shader = UsdShade.Shader.Define(
+        self.stage, mtl_path.AppendPath("Image_Texture")
+    )
+    uvmap_shader = UsdShade.Shader.Define(
+        self.stage, mtl_path.AppendPath("uvmap")
+    )
+
+    # setting the bsdf shader attributes
+    bsdf_shader.CreateIdAttr("UsdPreviewSurface")
+    bsdf_shader.CreateInput(
+        "diffuseColor", Sdf.ValueTypeNames.Color3f
+    ).ConnectToSource(image_shader.ConnectableAPI(), "rgb")
+    bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
+        float(self.rgba[-1])
+    )
+    bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
+        self.geom.shininess
+    )
+    bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
+        1.0 - self.geom.shininess
+    )
+
+    mtl.CreateSurfaceOutput().ConnectToSource(
+        bsdf_shader.ConnectableAPI(), "surface"
+    )
+
+    # setting the image texture attributes
+    image_shader.CreateIdAttr("UsdUVTexture")
+    image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(
+        self.texture_file
+    )
+    image_shader.CreateInput(
+        "sourceColorSpace", Sdf.ValueTypeNames.Token
+    ).Set("sRGB")
+    image_shader.CreateInput("wrapS", Sdf.ValueTypeNames.Token).Set("repeat")
+    image_shader.CreateInput("wrapT", Sdf.ValueTypeNames.Token).Set("repeat")
+    image_shader.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(
+        uvmap_shader.ConnectableAPI(), "result"
+    )
+    image_shader.CreateOutput("rgb", Sdf.ValueTypeNames.Float3)
+
+    # setting uvmap shader attributes
+    uvmap_shader.CreateIdAttr("UsdPrimvarReader_float2")
+    uvmap_shader.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("UVMap")
+    uvmap_shader.CreateOutput("results", Sdf.ValueTypeNames.Float2)
+
+    mtl.CreateSurfaceOutput().ConnectToSource(
+        bsdf_shader.ConnectableAPI(), "surface"
+    )
+
+    usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
+    UsdShade.MaterialBindingAPI(usd_mesh).Bind(mtl)
+
+  def attach_solid_material(self, usd_mesh):
+    """Attaches an solid texture to a material for a USD object"""
+    mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}")
+    mtl = UsdShade.Material.Define(self.stage, mtl_path)
+
+    bsdf_shader = UsdShade.Shader.Define(
+        self.stage, mtl_path.AppendPath("Principled_BSDF")
+    )
+
+    # settings the bsdf shader attributes
+    bsdf_shader.CreateIdAttr("UsdPreviewSurface")
+
+    bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(
+        tuple(self.rgba[0:3])
+    )
+    bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
+        float(self.rgba[-1])
+    )
+    bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
+        self.geom.shininess
+    )
+    bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
+        1.0 - self.geom.shininess
+    )
+
+    mtl.CreateSurfaceOutput().ConnectToSource(
+        bsdf_shader.ConnectableAPI(), "surface"
+    )
+
+    usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
+    UsdShade.MaterialBindingAPI(usd_mesh).Bind(mtl)
+
+  def _set_refinement_properties(self, usd_prim, scheme="none"):
+    usd_prim.GetAttribute("subdivisionScheme").Set(scheme)
+
+  def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int, scale: Optional[np.ndarray] = None):
+    """Updates the position and orientation of an object in the scene for a given frame"""
+    transformation_mat = utils_component.create_transform_matrix(
+        rotation_matrix=mat, translation_vector=pos
+    ).T
+    self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), 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
+
+    if scale is not None:
+      self.update_scale(scale, frame)
+
+  def update_visibility(self, visible: bool, frame: int):
+    """Updates the visibility of an object in a scene for a given frame"""
+    visibility_setting = "inherited" if visible else "invisible"
+    self.usd_xform.GetVisibilityAttr().Set(visibility_setting, frame)
+
+  def update_scale(self, scale: np.ndarray, frame: int):
+    """Updates the scale of an object in the scene for a given frame"""
+    self.scale_op.Set(Gf.Vec3f(scale.tolist()), frame)
+
+class USDMesh(USDObject):
+
+  def __init__(
+      self,
+      stage: Usd.Stage,
+      model: mujoco.MjModel,
+      geom: mujoco.MjvGeom,
+      obj_name: str,
+      dataid: int,
+      rgba: np.ndarray = np.array([1, 1, 1, 1]),
+      texture_file: Optional[str] = None,
+  ):
+    super().__init__(stage,
+                     geom,
+                     obj_name,
+                     rgba,
+                     texture_file)
+
+    self.model = model
+    self.dataid = dataid
+
+    mesh_path = f"{self.xform_path}/Mesh_{obj_name}"
+    self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
+    self.usd_prim = stage.GetPrimAtPath(mesh_path)
+
+    # setting mesh structure properties
+    mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
+    self.usd_mesh.GetPointsAttr().Set(mesh_vert)
+    self.usd_mesh.GetFaceVertexCountsAttr().Set(
+        [3 for _ in range(mesh_facenum)]
+    )
+    self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
+
+    # setting mesh uv properties
+    mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
+    self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
+        "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
+    )
+    self.texcoords.Set(mesh_texcoord)
+    self.texcoords.SetIndices(Vt.IntArray(mesh_facetexcoord.tolist()))
+
+    if self.texture_file:
+      self.attach_image_material(self.usd_mesh)
+    else:
+      self.attach_solid_material(self.usd_mesh)
+
+  def get_facetexcoord_ranges(self, nmesh, arr):
+    facetexcoords_ranges = [0]
+    running_sum = 0
+    for i in range(nmesh):
+      running_sum += arr[i] * 3
+      facetexcoords_ranges.append(running_sum)
+    return facetexcoords_ranges
+
+  def _get_uv_geometry(self):
+    mesh_texcoord_adr_from = self.model.mesh_texcoordadr[self.dataid]
+    mesh_texcoord_adr_to = (
+        self.model.mesh_texcoordadr[self.dataid + 1]
+        if self.dataid < self.model.nmesh - 1
+        else len(self.model.mesh_texcoord)
+    )
+    mesh_texcoord = self.model.mesh_texcoord[
+        mesh_texcoord_adr_from:mesh_texcoord_adr_to
+    ]
+
+    mesh_facetexcoord_ranges = self.get_facetexcoord_ranges(
+        self.model.nmesh, self.model.mesh_facenum
+    )
+
+    mesh_facetexcoord = self.model.mesh_facetexcoord.flatten()
+    mesh_facetexcoord = mesh_facetexcoord[
+        mesh_facetexcoord_ranges[self.dataid] : mesh_facetexcoord_ranges[
+            self.dataid + 1
+        ]
+    ]
+
+    mesh_facetexcoord[mesh_facetexcoord == len(mesh_texcoord)] = 0
+    return mesh_texcoord, mesh_facetexcoord
+
+  def _get_mesh_geometry(self):
+    mesh_vert_adr_from = self.model.mesh_vertadr[self.dataid]
+    mesh_vert_adr_to = (
+        self.model.mesh_vertadr[self.dataid + 1]
+        if self.dataid < self.model.nmesh - 1
+        else len(self.model.mesh_vert)
+    )
+    mesh_vert = self.model.mesh_vert[mesh_vert_adr_from:mesh_vert_adr_to]
+
+    mesh_face_adr_from = self.model.mesh_faceadr[self.dataid]
+    mesh_face_adr_to = (
+        self.model.mesh_faceadr[self.dataid + 1]
+        if self.dataid < self.model.nmesh - 1
+        else len(self.model.mesh_face)
+    )
+    mesh_face = self.model.mesh_face[mesh_face_adr_from:mesh_face_adr_to]
+    mesh_facenum = self.model.mesh_facenum[self.dataid]
+    return mesh_vert, mesh_face, mesh_facenum
+
+class USDPrimitiveMesh(USDObject):
+
+  def __init__(
+      self,
+      mesh_config: 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__(stage,
+                     geom,
+                     obj_name,
+                     rgba,
+                     texture_file)
+
+    self.mesh_config = mesh_config
+    self.prim_mesh = self.generate_primitive_mesh()
+
+    mesh_path = f"{self.xform_path}/Mesh_{obj_name}"
+    self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
+    self.usd_prim = stage.GetPrimAtPath(mesh_path)
+
+    mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
+    self.usd_mesh.GetPointsAttr().Set(mesh_vert)
+    self.usd_mesh.GetFaceVertexCountsAttr().Set(
+        [3 for _ in range(mesh_facenum)]
+    )
+    self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
+
+    # setting mesh uv properties
+    mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
+    self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
+        "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
+    )
+
+    self.texcoords.Set(mesh_texcoord)
+    self.texcoords.SetIndices(Vt.IntArray([i for i in range(mesh_facenum * 3)]))
+
+    self._set_refinement_properties(self.usd_prim)
+
+    if self.texture_file:
+      self.attach_image_material(self.usd_mesh)
+    else:
+      self.attach_solid_material(self.usd_mesh)
+
+  def generate_primitive_mesh(self):
+    _, prim_mesh = shapes_component.mesh_generator(self.mesh_config)
+    prim_mesh.translate(-prim_mesh.get_center())
+    return prim_mesh
+
+  def _get_uv_geometry(self):
+
+    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:
+      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
+
+    return mesh_texcoord, mesh_facetexcoord.flatten()
+
+  def _get_mesh_geometry(self):
+    assert self.prim_mesh
+
+    # get mesh geometry from the open3d mesh model
+    mesh_vert = np.asarray(self.prim_mesh.vertices)
+    mesh_face = np.asarray(self.prim_mesh.triangles)
+
+    return mesh_vert, mesh_face, len(mesh_face)
+
+class USDTendon(USDObject):
+  
+  def __init__(
+      self,
+      mesh_config: 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__(stage,
+                     geom,
+                     obj_name,
+                     rgba,
+                     texture_file)
+
+    self.mesh_config = mesh_config
+    self.tendon_parts = self.generate_primitive_mesh()
+    self.usd_refs = defaultdict(dict)
+
+    for name, mesh in self.tendon_parts.items():
+      part_xform_path = f"{self.xform_path}/Mesh_Xform_{name}"
+      mesh_path = f"{part_xform_path}/Mesh_{obj_name}"
+      self.usd_refs[name]["usd_xform"] = UsdGeom.Xform.Define(stage, part_xform_path)
+      self.usd_refs[name]["usd_mesh"] = UsdGeom.Mesh.Define(stage, mesh_path)
+      self.usd_refs[name]["usd_prim"] = stage.GetPrimAtPath(mesh_path)
+      # adding ops for each of the part xforms
+      self.usd_refs[name]["translate_op"] = self.usd_refs[name]["usd_xform"].AddTranslateOp()
+      self.usd_refs[name]["scale_op"] = self.usd_refs[name]["usd_xform"].AddScaleOp()
+
+    # setting mesh geometry properties for each of the parts in the tendon
+    part_geometries = self._get_mesh_geometry()
+    for name, part_geometry in part_geometries.items():
+      self.usd_refs[name]["usd_mesh"].GetPointsAttr().Set(part_geometry["mesh_vert"])
+      self.usd_refs[name]["usd_mesh"].GetFaceVertexCountsAttr().Set(
+        [3 for _ in range(part_geometry["mesh_facenum"])]
+      )
+      self.usd_refs[name]["usd_mesh"].GetFaceVertexIndicesAttr().Set(part_geometry["mesh_face"])
+
+    # setting uv properties for each of the parts in the tendon
+    part_uv_geometries = self._get_uv_geometry()
+    for name, part_uv_geometry in part_uv_geometries.items():
+      self.texcoords = UsdGeom.PrimvarsAPI(self.usd_refs[name]["usd_mesh"]).CreatePrimvar(
+          "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
+      )
+      self.texcoords.Set(part_uv_geometry["mesh_texcoord"])
+      self.texcoords.SetIndices(Vt.IntArray([i for i in range(part_geometry["mesh_facenum"] * 3)]))
+    
+    for name in self.usd_refs.keys(): 
+      self._set_refinement_properties(self.usd_refs[name]["usd_prim"])
+      if self.texture_file:
+        self.attach_image_material(self.usd_refs[name]["usd_mesh"])
+      else:
+        self.attach_solid_material(self.usd_refs[name]["usd_mesh"])
+    
+  def generate_primitive_mesh(self):
+    mesh_parts = {}
+    for part_config in self.mesh_config:
+      mesh_name, prim_mesh = shapes_component.mesh_generator(part_config)
+      prim_mesh.translate(-prim_mesh.get_center())
+      mesh_parts[mesh_name] = prim_mesh
+    return mesh_parts
+  
+  def _get_uv_geometry(self):
+    part_uv_geometries = defaultdict(dict)
+    for name, mesh in self.tendon_parts.items():
+      x_scale, y_scale = self.geom.texrepeat
+      mesh_texcoord = np.array(mesh.triangle_uvs)
+      mesh_facetexcoord = np.asarray(mesh.triangles)
+      part_uv_geometries[name] = {
+        "mesh_texcoord": mesh_texcoord,
+        "mesh_facetexcoord": mesh_facetexcoord
+      }
+    return part_uv_geometries
+
+  def _get_mesh_geometry(self):
+    part_geometries = defaultdict(dict)
+    for name, mesh in self.tendon_parts.items():
+      # get mesh geometry from the open3d mesh model
+      mesh_vert = np.asarray(mesh.vertices)
+      mesh_face = np.asarray(mesh.triangles)
+      part_geometries[name] = {
+        "mesh_vert": mesh_vert,
+        "mesh_face": mesh_face,
+        "mesh_facenum": len(mesh_face)
+      }
+    return part_geometries
+
+  def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int, scale: Optional[np.ndarray] = None):
+    """Updates the position and orientation of an object in the scene for a given frame"""
+    super().update(pos, mat, visible, frame, scale)
+    for name in self.tendon_parts.keys():
+      if "left" in name:
+        translate = [0, 0, -scale[2] - (scale[0] / 2)]
+        self.usd_refs[name]["translate_op"].Set(Gf.Vec3f(translate), frame)
+      elif "right" in name:
+        translate = [0, 0, scale[2] + (scale[0] / 2)]
+        self.usd_refs[name]["translate_op"].Set(Gf.Vec3f(translate), frame)
+
+  def update_scale(self, scale: np.ndarray, frame: int):
+    """Updates the scale of the tendon"""
+    for name in self.tendon_parts.keys():
+      if "cylinder" in name:
+        self.usd_refs[name]["scale_op"].Set(Gf.Vec3f(scale.tolist()), frame)
+      else:
+        hemisphere_scale = scale.tolist()
+        hemisphere_scale[2] = hemisphere_scale[0]
+        self.usd_refs[name]["scale_op"].Set(Gf.Vec3f(hemisphere_scale), frame)
+    
+    

python/mujoco/usd/shapes.py

@@ -27,14 +27,28 @@ def create_hemisphere(
 
     return mesh
 
+def decouple_config(config: dict):
+  decoupled_config = []
+  for key, value in config.items():
+    if key == "name":
+      continue
+    decoupled_config.append({
+      "parent_name": config["name"],
+      "name": config["name"] + "_" + key,
+      key: value.copy()
+    })
+
+  return decoupled_config
+
 def mesh_config_generator(
     name: str,
     geom_type: mujoco.mjtGeom,
-    size: np.ndarray
+    size: np.ndarray,
+    decouple: bool = False
 ):
 
     if geom_type == mujoco.mjtGeom.mjGEOM_PLANE:
-      return {
+      config = {
         "name": name,
         "box": {
           "width": size[0] * 2 if size[0] > 0 else 100,
@@ -44,7 +58,7 @@ def mesh_config_generator(
         }
       }
     elif geom_type == mujoco.mjtGeom.mjGEOM_SPHERE:
-      return {
+      config = {
         "name": name,
         "sphere": {
           "radius": float(size[0])
@@ -52,7 +66,7 @@ def mesh_config_generator(
       }
     elif geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE:
       cylinder = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_CYLINDER, size)
-      return {
+      config = {
         "name": name,
         "cylinder": cylinder["cylinder"],
         "left_hemisphere": {
@@ -74,12 +88,12 @@ def mesh_config_generator(
       sphere["sphere"]["transform"] = {
         "scale": tuple(size)
       }
-      return {
+      config = {
         "name": name,
         "sphere": sphere["sphere"],
       }
     elif geom_type == mujoco.mjtGeom.mjGEOM_CYLINDER:
-      return {
+      config = {
         "name": name,
         "cylinder": {
           "radius": size[0],
@@ -87,7 +101,7 @@ def mesh_config_generator(
         }
       }
     elif geom_type == mujoco.mjtGeom.mjGEOM_BOX:
-      return {
+      config = {
         "name": name,
         "box": {
           "width": size[0] * 2,
@@ -98,6 +112,10 @@ def mesh_config_generator(
     else:
       raise NotImplemented(f"{geom_type} primitive geom type not implemented with USD integration")
 
+    if decouple:
+      config = decouple_config(config)
+
+    return config
 
 def mesh_generator(
     mesh_config: dict,
@@ -110,7 +128,7 @@ def mesh_generator(
 
     for shape, config in mesh_config.items():
 
-      if shape == "name":
+      if "name" in shape:
         continue
 
       if "box" in shape:
@@ -140,9 +158,9 @@ def mesh_generator(
         )
 
       if "transform" in config:
-
         if "rotate" in config["transform"]:
-          R = mesh.get_rotation_matrix_from_xyz(config["transform"]["rotate"])
+          tmp = o3d.geometry.TriangleMesh.create_coordinate_frame() # TODO: add rotation matrix code gen to utils.py
+          R = tmp.get_rotation_matrix_from_xyz(config["transform"]["rotate"])
           prim_mesh.rotate(R, center=(0, 0, 0))
         if "scale" in config["transform"]:
           prim_mesh.vertices = o3d.utility.Vector3dVector(

python/mujoco/usd/temp.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from abs import ABC
 from typing import List, Optional, Tuple
 
 import mujoco
@@ -34,7 +35,6 @@ from pxr import UsdLux
 from pxr import UsdShade
 from pxr import Vt
 
-
 class USDMesh:
 
   def __init__(
@@ -265,7 +265,7 @@ class USDPrimitiveMesh:
     self.prim_mesh = None
     self.transform_op = Gf.Matrix4d(1.)
 
-    _, self.prim_mesh = shapes_component.mesh_generator(mesh_config)
+    self.prim_mesh = self.generate_primitive_mesh()
 
     xform_path = f"/World/{self.obj_name}_Xform"
     mesh_path = f"{xform_path}/{obj_name}"
@@ -273,8 +273,6 @@ class USDPrimitiveMesh:
     self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
     self.usd_prim = stage.GetPrimAtPath(mesh_path)
 
-    self.prim_mesh.translate(-self.prim_mesh.get_center())
-
     mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
     self.usd_mesh.GetPointsAttr().Set(mesh_vert)
     self.usd_mesh.GetFaceVertexCountsAttr().Set(
@@ -302,6 +300,11 @@ class USDPrimitiveMesh:
 
     self.last_visible_frame = -2
 
+  def generate_primitive_mesh(self):
+    _, prim_mesh = shapes_component.mesh_generator(self.mesh_config)
+    prim_mesh.translate(-prim_mesh.get_center())
+    return prim_mesh
+
   def _set_refinement_properties(self):
     self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
 
@@ -417,7 +420,7 @@ class USDPrimitiveMesh:
     self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
     UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
 
-  def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
+  def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int, ):
     transformation_mat = utils_component.create_transform_matrix(
         rotation_matrix=mat, translation_vector=pos
     ).T
@@ -438,29 +441,12 @@ class USDPrimitiveMesh:
   def update_scale(self, scale: np.ndarray, frame: int):
     self.scale_op.Set(Gf.Vec3f(scale.tolist()), frame)
 
-class USDTendon(USDPrimitiveMesh):
-
+class USDTendon:
+  
   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,
+      self
   ):
-    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)
+    pass
 
 class USDSphereLight: