combining shapes into one class

python/mujoco/usd/__init__.py

@@ -1,3 +1,4 @@
 from .exporter import *
 from .component import *
-from .utils import *
+from .utils import *
+from .shapes import *

python/mujoco/usd/component.py

@@ -15,7 +15,12 @@
 from typing import List, Optional, Tuple
 
 import mujoco
-import mujoco.usd.utils
+
+# import mujoco.usd.utils
+# 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.
@@ -217,7 +222,7 @@ class USDMesh:
     UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
 
   def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
-    transformation_mat = mujoco.usd.utils.create_transform_matrix(
+    transformation_mat = utils_component.create_transform_matrix(
         rotation_matrix=mat, translation_vector=pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
@@ -234,12 +239,14 @@ class 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,
   ):
+    self.mesh_config = mesh_config
     self.stage = stage
     self.geom = geom
     self.obj_name = obj_name
@@ -251,6 +258,40 @@ class USDPrimitiveMesh:
     self.prim_mesh = None
     self.transform_op = Gf.Matrix4d(1.)
 
+    _, self.prim_mesh = shapes_component.mesh_generator(mesh_config)
+
+    xform_path = f"/World/{self.obj_name}_Xform"
+    mesh_path = f"{xform_path}/{obj_name}"
+    self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
+    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(
+        [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()
+
+    # setting attributes for the shape
+    self._attach_material()
+
+    # defining ops required by update function
+    self.transform_op = self.usd_xform.AddTransformOp()
+
   def _set_refinement_properties(self):
     self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
 
@@ -365,7 +406,7 @@ class USDPrimitiveMesh:
     UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
 
   def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
-    transformation_mat = mujoco.usd.utils.create_transform_matrix(
+    transformation_mat = utils_component.create_transform_matrix(
         rotation_matrix=mat, translation_vector=pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
@@ -377,284 +418,6 @@ class USDPrimitiveMesh:
     else:
       self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
 
-
-class USDCapsule(USDPrimitive):
-
-  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,
-  ):
-
-    super().__init__(stage, geom, obj_name, rgba, texture_file)
-
-    xform_path = f"/World/Capsule_Xform_{obj_name}"
-    capsule_path = f"{xform_path}/Capsule_{obj_name}"
-    self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
-    self.usd_primitive_shape = UsdGeom.Capsule.Define(stage, capsule_path)
-    self.usd_prim = stage.GetPrimAtPath(capsule_path)
-
-    # defining ops required by update function
-    self.transform_op = self.usd_xform.AddTransformOp()
-    self.scale_op = self.usd_xform.AddScaleOp()
-
-    # setting attributes for the shape
-    self._set_size_attributes()
-    self._attach_material()
-
-    # self._set_refinement_properties()
-
-  def _set_size_attributes(self):
-    self.usd_primitive_shape.GetRadiusAttr().Set(float(self.geom.size[0]))
-    self.usd_primitive_shape.GetHeightAttr().Set(
-        float(self.geom.size[2] * 2)
-    )  # mujoco gives the half length
-
-
-class USDEllipsoid(USDPrimitive):
-
-  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,
-  ):
-
-    super().__init__(stage, geom, obj_name, rgba, texture_file)
-
-    xform_path = f"/World/Ellipsoid_Xform_{obj_name}"
-    ellipsoid_path = f"{xform_path}/Ellipsoid_{obj_name}"
-    self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
-    self.usd_primitive_shape = UsdGeom.Sphere.Define(stage, ellipsoid_path)
-    self.usd_prim = stage.GetPrimAtPath(ellipsoid_path)
-
-    # defining ops required by update function
-    self.transform_op = self.usd_xform.AddTransformOp()
-    self.scale_op = self.usd_xform.AddScaleOp()
-
-    # setting attributes for the shape
-    self._set_size_attributes()
-    self._attach_material()
-
-    # self._set_refinement_properties()
-
-  def _set_size_attributes(self):
-    self.scale_op.Set(Gf.Vec3d(self.geom.size.tolist()))
-
-
-class USDCubeMesh(USDPrimitiveMesh):
-
-  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,
-  ):
-
-    super().__init__(stage, geom, obj_name, rgba, texture_file)
-
-    xform_path = f"/World/CubeMesh_Xform_{obj_name}"
-    mesh_path = f"{xform_path}/CubeMesh_{obj_name}"
-    self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
-    self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
-    self.usd_prim = stage.GetPrimAtPath(mesh_path)
-
-    self.prim_mesh = o3d.geometry.TriangleMesh.create_box(
-        width=self.geom.size[0] * 2,
-        height=self.geom.size[1] * 2,
-        depth=self.geom.size[2] * 2,
-    )
-
-    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(
-        [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()
-
-    # setting attributes for the shape
-    self._attach_material()
-
-    # defining ops required by update function
-    self.transform_op = self.usd_xform.AddTransformOp()
-
-
-class USDSphereMesh(USDPrimitiveMesh):
-
-  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,
-  ):
-
-    super().__init__(stage, geom, obj_name, rgba, texture_file)
-
-    xform_path = f"/World/SphereMesh_Xform_{obj_name}"
-    mesh_path = f"{xform_path}/SphereMesh_{obj_name}"
-    self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
-    self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
-    self.usd_prim = stage.GetPrimAtPath(mesh_path)
-
-    self.prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
-        radius=float(self.geom.size[0]), create_uv_map=True
-    )
-
-    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(
-        [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()
-
-    # setting attributes for the shape
-    self._attach_material()
-
-    # defining ops required by update function
-    self.transform_op = self.usd_xform.AddTransformOp()
-
-
-class USDCylinderMesh(USDPrimitiveMesh):
-
-  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,
-  ):
-
-    super().__init__(stage, geom, obj_name, rgba, texture_file)
-
-    xform_path = f"/World/CylinderMesh_Xform_{obj_name}"
-    mesh_path = f"{xform_path}/CylinderMesh_{obj_name}"
-    self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
-    self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
-    self.usd_prim = stage.GetPrimAtPath(mesh_path)
-
-    self.prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
-        radius=self.geom.size[0],
-        height=self.geom.size[2] * 2,
-        create_uv_map=True,
-    )
-
-    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(
-        [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()
-
-    # setting attributes for the shape
-    self._attach_material()
-
-    # defining ops required by update function
-    self.transform_op = self.usd_xform.AddTransformOp()
-
-
-class USDPlaneMesh(USDPrimitiveMesh):
-
-  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,
-  ):
-
-    super().__init__(stage, geom, obj_name, rgba, texture_file)
-
-    xform_path = f"/World/Plane_Xform_{obj_name}"
-    plane_path = f"{xform_path}/PlaneMesh_{obj_name}"
-    self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
-    self.usd_mesh = UsdGeom.Mesh.Define(stage, plane_path)
-    self.usd_prim = stage.GetPrimAtPath(plane_path)
-
-    self.prim_mesh = o3d.geometry.TriangleMesh.create_box(
-        width=self.geom.size[0] * 2 if self.geom.size[0] > 0 else 100,
-        height=self.geom.size[1] * 2 if self.geom.size[1] > 0 else 100,
-        depth=0.001,
-        create_uv_map=True,
-        map_texture_to_each_face=True,
-    )
-
-    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(
-        [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()
-
-    # setting attributes for the shape
-    self._attach_material()
-
-    # defining ops required by update function
-    self.transform_op = self.usd_xform.AddTransformOp()
-
-
 class USDSphereLight:
 
   def __init__(
@@ -732,7 +495,7 @@ class USDCamera:
 
   def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
 
-    transformation_mat = mujoco.usd.utils.create_transform_matrix(
+    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/demo.py

@@ -1,19 +1,21 @@
+
 import mujoco
-from mujoco.usd import exporter
+# from mujoco.usd import exporter
+import exporter
 
 if __name__ == "__main__":
 
   # load a model to mujoco
-  m = mujoco.MjModel.from_xml_path("/Users/abhishek/Documents/research/mujoco/model/humanoid/humanoid.xml")
+  m = mujoco.MjModel.from_xml_path("/Users/abhishek/Documents/research/mujoco/model/car/car.xml")
   d = mujoco.MjData(m)
 
   # create an instance of the USDExporter
   exp = exporter.USDExporter(model=m)
 
-  mujoco.mj_step(m, d)
+  for i in range(100):
+    mujoco.mj_step(m, d)
+    exp.update_scene(d)
 
-  exp.update_scene(d)
-
-  exp.save_scene(filetype="usda")
+  exp.save_scene(filetype="usd")
 
 

python/mujoco/usd/exporter.py

@@ -15,7 +15,12 @@
 import os
 
 import mujoco
-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
 import termcolor
@@ -234,7 +239,6 @@ class USDExporter:
 
     texture_file = self.texture_files[geom.texid] if geom.texid != -1 else None
 
-    # handles meshes in scene
     if geom.type == mujoco.mjtGeom.mjGEOM_MESH:
       usd_geom = component_module.USDMesh(
           stage=self.stage,
@@ -245,56 +249,22 @@ class USDExporter:
           rgba=geom.rgba,
           texture_file=texture_file,
       )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_PLANE:
-      usd_geom = component_module.USDPlaneMesh(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_SPHERE:
-      usd_geom = component_module.USDSphereMesh(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_CAPSULE:
-      usd_geom = component_module.USDCapsule(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
-      usd_geom = component_module.USDEllipsoid(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_CYLINDER:
-      usd_geom = component_module.USDCylinderMesh(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_BOX:
-      usd_geom = component_module.USDCubeMesh(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
     else:
-      usd_geom = None
+
+      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_name2usd[geom_name] = usd_geom
 

python/mujoco/usd/shapes.py

@@ -0,0 +1,129 @@
+import copy
+import mujoco
+import pprint
+import numpy as np
+import open3d as o3d
+
+def mesh_config_generator(
+    name: str,
+    geom_type: mujoco.mjtGeom,
+    size: np.ndarray
+):
+
+    if geom_type == mujoco.mjtGeom.mjGEOM_PLANE:
+      return {
+        "name": name,
+        "box": {
+          "width": size[0] * 2 if size[0] > 0 else 100,
+          "height": size[1] * 2 if size[1] > 0 else 100,
+          "depth": 0.001,
+          "map_texture_to_each_face": True,
+        }
+      }
+    elif geom_type == mujoco.mjtGeom.mjGEOM_SPHERE:
+      return {
+        "name": name,
+        "sphere": {
+          "radius": float(size[0])
+        }
+      }
+    elif geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE:
+      cylinder = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_CYLINDER, size)
+      left_sphere = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_SPHERE, size)
+      right_sphere = copy.deepcopy(left_sphere)
+      left_sphere["sphere"]["transform"] = {
+        "translate": (0, 0, -size[2])
+      }
+      right_sphere["sphere"]["transform"] = {
+        "translate": (0, 0, size[2])
+      } 
+      return {
+        "name": name,
+        "cylinder": cylinder["cylinder"],
+        "left_sphere": left_sphere["sphere"],
+        "right_sphere": right_sphere["sphere"],
+      }
+    elif geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
+      sphere = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_SPHERE, [1.0])
+      sphere["sphere"]["transform"] = {
+        "scale": tuple(size)
+      }
+      return {
+        "name": name,
+        "sphere": sphere["sphere"],
+      }
+    elif geom_type == mujoco.mjtGeom.mjGEOM_CYLINDER:
+      return {
+        "name": name,
+        "cylinder": {
+          "radius": size[0],
+          "height": size[2] * 2,
+        }
+      }
+    elif geom_type == mujoco.mjtGeom.mjGEOM_BOX:
+      return {
+        "name": name,
+        "box": {
+          "width": size[0] * 2,
+          "height": size[1] * 2,
+          "depth": size[2] * 2,
+        }
+      }
+    else:
+      raise NotImplemented(f"{geom_type} primitive geom type not implemented with USD integration")
+
+
+def mesh_generator(
+    mesh_config: dict
+):
+
+    assert "name" in mesh_config
+
+    mesh = None
+
+    for shape, config in mesh_config.items():
+
+      if shape == "name":
+        continue
+
+      if "box" in shape:
+        prim_mesh = o3d.geometry.TriangleMesh.create_box(
+            width=mesh_config[shape]["width"],
+            height=mesh_config[shape]["height"],
+            depth=mesh_config[shape]["depth"],
+            create_uv_map=True,
+            map_texture_to_each_face=True,
+        )
+      elif "sphere" in shape:
+        prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
+            radius=mesh_config[shape]["radius"],
+            create_uv_map=True
+        )
+      elif "cylinder" in shape:
+        prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
+            radius=mesh_config[shape]["radius"],
+            height=mesh_config[shape]["height"],
+            create_uv_map=True,
+        )
+
+      if "transform" in config:
+        for transform, val in config["transform"].items():
+          if transform == "translate":
+            prim_mesh.translate(val)
+          if transform == "scale":
+            prim_mesh.vertices = o3d.utility.Vector3dVector(
+              np.asarray(prim_mesh.vertices) * np.array(val))
+
+      if not mesh:
+        mesh = prim_mesh
+      else:
+        mesh += prim_mesh
+
+    return mesh_config["name"], mesh
+      
+
+
+
+
+    
+