Merge pull request #2 from awesome-aj0123/demo_script

Demo script
google-deepmind · Apr 2, 2024 · 09bc73e · 09bc73e
2 parents d7f131e + a4b5355
commit 09bc73e
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 123 deletions.
diff --git a/.gitignore b/.gitignore
@@ -18,6 +18,7 @@
 *.egg-info/
 build/
 build_cmake/
+python/dist/
 
 # Exclude macOS folder attributes
 .DS_Store
@@ -28,4 +29,4 @@ Info.framework.plist
 # Python byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
-*$py.class
+*$py.class
diff --git a/python/mujoco/usd/__init__.py b/python/mujoco/usd/__init__.py
@@ -0,0 +1,3 @@
+from .exporter import *
+from .component import *
+from .utils import *
diff --git a/python/mujoco/usd/component.py b/python/mujoco/usd/component.py
@@ -378,123 +378,6 @@ def update_visibility(self, visible: bool, frame: int):
       self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
 
 
-class 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,
-  ):
-    self.stage = stage
-    self.geom = geom
-    self.obj_name = obj_name
-    self.rgba = rgba
-    self.texture_file = texture_file
-
-    self.usd_prim = Usd.Prim()
-    self.usd_primitive_shape = Usd.PrimitiveShape()
-    self.transform_op = Usd.TransformOp()
-
-  def _set_refinement_properties(self):
-    self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
-
-  def _attach_material(self):
-    mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}")
-    mtl = UsdShade.Material.Define(self.stage, mtl_path)
-    if self.texture_file:
-      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")
-      )
-
-      # settings 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"
-      )
-
-      self.usd_primitive_shape.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
-      UsdShade.MaterialBindingAPI(self.usd_primitive_shape).Bind(mtl)
-
-      # 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("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)
-    else:
-      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[: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"
-    )
-
-    self.usd_primitive_shape.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
-    UsdShade.MaterialBindingAPI(self.usd_primitive_shape).Bind(mtl)
-
-  def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
-    transformation_mat = mujoco.usd_util.create_transform_matrix(
-        rotation_matrix=mat, translation_vector=pos
-    ).T
-    self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
-    self.update_visibility(visible, 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)
-
-
 class USDCapsule(USDPrimitive):
 
   def __init__(
@@ -522,7 +405,7 @@ def __init__(
     self._set_size_attributes()
     self._attach_material()
 
-    self._set_refinement_properties()
+    # self._set_refinement_properties()
 
   def _set_size_attributes(self):
     self.usd_primitive_shape.GetRadiusAttr().Set(float(self.geom.size[0]))
@@ -558,7 +441,7 @@ def __init__(
     self._set_size_attributes()
     self._attach_material()
 
-    self._set_refinement_properties()
+    # self._set_refinement_properties()
 
   def _set_size_attributes(self):
     self.scale_op.Set(Gf.Vec3d(self.geom.size.tolist()))
@@ -849,7 +732,7 @@ def __init__(self, stage: Usd.Stage, obj_name: str):
 
   def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
 
-    transformation_mat = mujoco.usd_util.create_transform_matrix(
+    transformation_mat = mujoco.usd.utils.create_transform_matrix(
         rotation_matrix=cam_mat, translation_vector=cam_pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
diff --git a/python/mujoco/usd/demo.py b/python/mujoco/usd/demo.py
@@ -0,0 +1,19 @@
+import mujoco
+from mujoco.usd 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")
+  d = mujoco.MjData(m)
+
+  # create an instance of the USDExporter
+  exp = exporter.USDExporter(model=m)
+
+  mujoco.mj_step(m, d)
+
+  exp.update_scene(d)
+
+  exp.save_scene(filetype="usda")
+
+
diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py
@@ -395,7 +395,7 @@ def add_light(
       intensity: int,
       radius: Optional[float] = 1.0,
       color: Optional[np.ndarray] = np.array([0.3, 0.3, 0.3]),
-      objid: Optional[int] = 1,
+      obj_name: Optional[str] = "light_1",
       light_type: Optional[str] = "sphere",
   ):
 
@@ -413,7 +413,7 @@ def add_camera(
       self,
       pos: List[float],
       rotation_xyz: List[float],
-      objid: Optional[int] = 1,
+      obj_name: Optional[str] = "camera_1",
   ):
     new_camera = component_module.USDCamera(
         stage=self.stage, obj_name=str(objid))

diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py