addressing PR comments

python/mujoco/__init__.py

@@ -46,6 +46,9 @@ from mujoco._render import *
 from mujoco._structs import *
 from mujoco.gl_context import *
 from mujoco.renderer import Renderer
+from mujoco.usd.usd_exporter import USDExporter
+from mujoco.usd.usd_component import *
+from mujoco.usd.usd_utils import *
 
 HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco')
 PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin')

python/mujoco/usd/usd_component.py

@@ -1,12 +1,15 @@
+from typing import Optional, List, Tuple
+
 import numpy as np
 import open3d as o3d
-from usd_utils import *
-from mujoco import mjtGeom
 from pxr import Gf, Sdf, Vt
-from typing import Optional, List
-from mujoco import _structs, _enums
 from pxr import Usd, UsdGeom, UsdShade, UsdLux
 
+import mujoco
+from mujoco import mjtGeom
+from mujoco.usd.usd_utils import *
+from mujoco import _structs, _enums
+
 class USDMesh:
 
     def __init__(
@@ -16,7 +19,7 @@ class USDMesh:
         geom: _structs.MjvGeom,
         objid: int,
         dataid: int,
-        rgba: List[int] = [1,1,1,1],
+        rgba: Tuple[int] = (1,1,1,1),
         texture_file: Optional[str] = None
     ):
         """ Initializes a new USD mesh
@@ -58,7 +61,11 @@ class USDMesh:
         # defining ops required by update function
         self.transform_op = self.usd_xform.AddTransformOp()
 
-    def get_facetexcoord_ranges(self, nmesh, arr):
+    def get_facetexcoord_ranges(
+        self, 
+        nmesh, 
+        arr
+    ):
         facetexcoords_ranges = [0]
         running_sum = 0
         for i in range(nmesh):
@@ -68,7 +75,9 @@ class USDMesh:
 
     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_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)
@@ -76,16 +85,21 @@ class USDMesh:
         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):
-        # get mesh geometry structure from reading the mjModel
         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_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_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]
@@ -105,8 +119,8 @@ class USDMesh:
             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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(self.geom.shininess)
-            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+            bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(1.0 - self.geom.shininess)
 
             mtl.CreateSurfaceOutput().ConnectToSource(bsdf_shader.ConnectableAPI(), "surface")
 
@@ -131,10 +145,11 @@ class USDMesh:
 
             # settings the bsdf shader attributes
             bsdf_shader.CreateIdAttr("UsdPreviewSurface")
-            bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(tuple(self.rgba[:3]))
+
+            bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(tuple(color))
             bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(float(self.rgba[-1]))
-            bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(self.geom.shininess)
-            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+            bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(1.0 - self.geom.shininess)
 
         mtl.CreateSurfaceOutput().ConnectToSource(bsdf_shader.ConnectableAPI(), "surface")
 
@@ -145,19 +160,33 @@ class USDMesh:
         self,
         pos: np.array,
         mat: np.array,
+        visible: bool,
         frame: int
     ):
         transformation_mat = 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 USDPrimitiveMesh:
 
-    def __init__(self,
+    def __init__(
-                 stage: Usd.Stage,
+        self,
-                 geom: _structs.MjvGeom,
+        stage: Usd.Stage,
-                 objid: int,
+        geom: _structs.MjvGeom,
-                 rgba: List[int] = [1,1,1,1],
+        objid: int,
-                 texture_file: Optional[str] = None):
+        rgba: Tuple[int] = (1, 1, 1, 1),
+        texture_file: Optional[str] = None
+    ):
         self.stage = stage
         self.geom = geom
         self.objid = objid
@@ -166,6 +195,9 @@ class USDPrimitiveMesh:
 
         self.prim_mesh = None
 
+    def _set_refinement_properties(self):
+        self.usd_prim.GetAttribute('subdivisionScheme').Set("none")
+
     def _get_uv_geometry(self):    
 
         assert self.prim_mesh
@@ -175,8 +207,12 @@ class USDPrimitiveMesh:
         mesh_texcoord = np.array(self.prim_mesh.triangle_uvs)
         mesh_facetexcoord = np.asarray(self.prim_mesh.triangles)
 
-        mesh_texcoord[:, 0] *= x_scale
+        x_multiplier, y_multiplier = 1, 1
-        mesh_texcoord[:, 1] *= y_scale
+        if self.geom.texuniform:
+            x_multiplier, y_multiplier = self.geom.size[:2]
+
+        mesh_texcoord[:, 0] *= x_scale * x_multiplier
+        mesh_texcoord[:, 1] *= y_scale * y_multiplier
 
         return mesh_texcoord, mesh_facetexcoord.flatten()
 
@@ -194,7 +230,6 @@ class USDPrimitiveMesh:
         mtl_path = Sdf.Path(f"/World/_materials/Material_{self.objid}")
         mtl = UsdShade.Material.Define(self.stage, mtl_path)
         if self.texture_file:
-            # remove all code in this if block
             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"))
@@ -203,8 +238,8 @@ class USDPrimitiveMesh:
             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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(self.geom.shininess)
-            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+            bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(1.0 - self.geom.shininess)
 
             mtl.CreateSurfaceOutput().ConnectToSource(bsdf_shader.ConnectableAPI(), "surface")
 
@@ -231,8 +266,8 @@ class USDPrimitiveMesh:
             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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(self.geom.shininess)
-            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+            bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(1.0 - self.geom.shininess)
 
         mtl.CreateSurfaceOutput().ConnectToSource(bsdf_shader.ConnectableAPI(), "surface")
 
@@ -241,26 +276,43 @@ class USDPrimitiveMesh:
 
     def update(
         self,
-        pos: np.array,
+        pos: np.ndarray,
-        mat: np.array,
+        mat: np.ndarray,
+        visible: bool,
         frame: int
     ):
         transformation_mat = 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 USDPrimitive:
-    def __init__(self,
+    def __init__(
-                 stage: Usd.Stage,
+        self,
-                 geom: _structs.MjvGeom,
+        stage: Usd.Stage,
-                 objid: int,
+        geom: _structs.MjvGeom,
-                 rgba: List[int] = [1,1,1,1],
+        objid: int,
-                 texture_file: Optional[str] = None):
+        rgba: Tuple[int] = (1, 1, 1, 1),
+        texture_file: Optional[str] = None
+    ):
         self.stage = stage
         self.geom = geom
         self.objid = objid
         self.rgba = rgba
         self.texture_file = texture_file
     
+    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.objid}")
         mtl = UsdShade.Material.Define(self.stage, mtl_path)
@@ -273,8 +325,8 @@ class USDPrimitive:
             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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(self.geom.shininess)
-            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+            bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(1.0 - self.geom.shininess)
 
             mtl.CreateSurfaceOutput().ConnectToSource(bsdf_shader.ConnectableAPI(), "surface")
 
@@ -299,8 +351,8 @@ class USDPrimitive:
             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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(self.geom.shininess)
-            bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+            bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(1.0 - self.geom.shininess)
 
         mtl.CreateSurfaceOutput().ConnectToSource(bsdf_shader.ConnectableAPI(), "surface")
 
@@ -309,26 +361,42 @@ class USDPrimitive:
 
     def update(
         self,
-        pos: np.array,
+        pos: np.ndarray,
-        mat: np.array,
+        mat: np.ndarray,
+        visible: bool,
         frame: int
     ):
         transformation_mat = 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__(self,
+    def __init__(
-                 stage: Usd.Stage,
+        self,
-                 geom: _structs.MjvGeom,
+        stage: Usd.Stage,
-                 objid: int,
+        geom: _structs.MjvGeom,
-                 rgba: List[int] = [1,1,1,1],
+        objid: int,
-                 texture_file: Optional[str] = None):
+        rgba: Tuple[int] = (1, 1, 1, 1),
+        texture_file: Optional[str] = None
+    ):
         
-        super().__init__(stage,
+        super().__init__(
-                         geom,
+            stage,
-                         objid,
+            geom,
-                         rgba,
+            objid,
-                         texture_file)
+            rgba,
+            texture_file
+        )
         
         xform_path = f'/World/Capsule_Xform_{objid}'
         capsule_path = f'{xform_path}/Capsule_{objid}'
@@ -344,23 +412,29 @@ class USDCapsule(USDPrimitive):
         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,
+    def __init__(
-                 stage: Usd.Stage,
+        self,
-                 geom: _structs.MjvGeom,
+        stage: Usd.Stage,
-                 objid: int,
+        geom: _structs.MjvGeom,
-                 rgba: List[int] = [1,1,1,1],
+        objid: int,
-                 texture_file: Optional[str] = None):
+        rgba: Tuple[int] = (1, 1, 1, 1),
+        texture_file: Optional[str] = None
+    ):
         
-        super().__init__(stage,
+        super().__init__(
-                         geom,
+            stage,
-                         objid,
+            geom,
-                         rgba,
+            objid,
-                         texture_file)
+            rgba,
+            texture_file
+        )
         
         xform_path = f'/World/Ellipsoid_Xform_{objid}'
         ellipsoid_path = f'{xform_path}/Ellipsoid_{objid}'
@@ -376,22 +450,28 @@ class USDEllipsoid(USDPrimitive):
         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,
+    def __init__(
-                 stage: Usd.Stage,
+        self,
-                 geom: _structs.MjvGeom,
+        stage: Usd.Stage,
-                 objid: int,
+        geom: _structs.MjvGeom,
-                 rgba: List[int] = [1,1,1,1],
+        objid: int,
-                 texture_file: Optional[str] = None):
+        rgba: Tuple[int] = (1, 1, 1, 1),
+        texture_file: Optional[str] = None
+    ):
         
-        super().__init__(stage,
+        super().__init__(
-                         geom,
+            stage,
-                         objid,
+            geom,
-                         rgba,
+            objid,
-                         texture_file)
+            rgba,
+            texture_file
+        )
 
         xform_path = f'/World/CubeMesh_Xform_{objid}'
         mesh_path= f'{xform_path}/CubeMesh_{objid}'
@@ -399,11 +479,13 @@ class USDCubeMesh(USDPrimitiveMesh):
         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, 
+        self.prim_mesh = o3d.geometry.TriangleMesh.create_box(
-                                                              height=self.geom.size[1]*2,
+            width=self.geom.size[0]*2, 
-                                                              depth=self.geom.size[2]*2,
+            height=self.geom.size[1]*2,
-                                                              create_uv_map=True,
+            depth=self.geom.size[2]*2,
-                                                              map_texture_to_each_face=True)
+            create_uv_map=True,
+            map_texture_to_each_face=True
+        )
 
         self.prim_mesh.translate(-self.prim_mesh.get_center())
 
@@ -421,6 +503,8 @@ class USDCubeMesh(USDPrimitiveMesh):
         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()
 
@@ -428,18 +512,22 @@ class USDCubeMesh(USDPrimitiveMesh):
         self.transform_op = self.usd_xform.AddTransformOp()
 
 class USDSphereMesh(USDPrimitiveMesh):
-    def __init__(self,
+    def __init__(
-                 stage: Usd.Stage,
+        self,
-                 geom: _structs.MjvGeom,
+        stage: Usd.Stage,
-                 objid: int,
+        geom: _structs.MjvGeom,
-                 rgba: List[int] = [1,1,1,1],
+        objid: int,
-                 texture_file: Optional[str] = None):
+        rgba: Tuple[int] = (1, 1, 1, 1),
+        texture_file: Optional[str] = None
+    ):
         
-        super().__init__(stage,
+        super().__init__(
-                         geom,
+            stage,
-                         objid,
+            geom,
-                         rgba,
+            objid,
-                         texture_file)
+            rgba,
+            texture_file
+        )
 
         xform_path = f'/World/SphereMesh_Xform_{objid}'
         mesh_path= f'{xform_path}/SphereMesh_{objid}'
@@ -447,8 +535,10 @@ class USDSphereMesh(USDPrimitiveMesh):
         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]),
+        self.prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
-                                                                 create_uv_map=True)
+            radius=float(self.geom.size[0]),
+            create_uv_map=True
+        )
 
         self.prim_mesh.translate(-self.prim_mesh.get_center())
 
@@ -466,6 +556,8 @@ class USDSphereMesh(USDPrimitiveMesh):
         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()
 
@@ -473,18 +565,22 @@ class USDSphereMesh(USDPrimitiveMesh):
         self.transform_op = self.usd_xform.AddTransformOp()
 
 class USDCylinderMesh(USDPrimitiveMesh):
-    def __init__(self,
+    def __init__(
-                 stage: Usd.Stage,
+        self,
-                 geom: _structs.MjvGeom,
+        stage: Usd.Stage,
-                 objid: int,
+        geom: _structs.MjvGeom,
-                 rgba: List[int] = [1,1,1,1],
+        objid: int,
-                 texture_file: Optional[str] = None):
+        rgba: Tuple[int] = (1, 1, 1, 1),
+        texture_file: Optional[str] = None
+    ):
         
-        super().__init__(stage,
+        super().__init__(
-                         geom,
+            stage,
-                         objid,
+            geom,
-                         rgba,
+            objid,
-                         texture_file)
+            rgba,
+            texture_file
+        )
 
         xform_path = f'/World/CylinderMesh_Xform_{objid}'
         mesh_path= f'{xform_path}/CylinderMesh_{objid}'
@@ -492,9 +588,11 @@ class USDCylinderMesh(USDPrimitiveMesh):
         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],
+        self.prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
-                                                                   height=self.geom.size[2]*2,
+            radius=self.geom.size[0],
-                                                                   create_uv_map=True)
+            height=self.geom.size[2]*2,
+            create_uv_map=True
+        )
 
         self.prim_mesh.translate(-self.prim_mesh.get_center())
 
@@ -512,6 +610,8 @@ class USDCylinderMesh(USDPrimitiveMesh):
         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()
 
@@ -519,18 +619,22 @@ class USDCylinderMesh(USDPrimitiveMesh):
         self.transform_op = self.usd_xform.AddTransformOp()
 
 class USDPlaneMesh(USDPrimitiveMesh):
-    def __init__(self,
+    def __init__(
-                 stage: Usd.Stage,
+        self,
-                 geom: _structs.MjvGeom,
+        stage: Usd.Stage,
-                 objid: int,
+        geom: _structs.MjvGeom,
-                 rgba: List[int] = [1,1,1,1],
+        objid: int,
-                 texture_file: Optional[str] = None):
+        rgba: Tuple[int] = (1, 1, 1, 1),
+        texture_file: Optional[str] = None
+    ):
 
-        super().__init__(stage,
+        super().__init__(
-                         geom,
+            stage,
-                         objid,
+            geom,
-                         rgba,
+            objid,
-                         texture_file)
+            rgba,
+            texture_file
+        )
         
         xform_path = f'/World/Plane_Xform_{objid}'
         plane_path = f'{xform_path}/PlaneMesh_{objid}'
@@ -538,11 +642,13 @@ class USDPlaneMesh(USDPrimitiveMesh):
         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, 
+        self.prim_mesh = o3d.geometry.TriangleMesh.create_box(
-                                                              height=self.geom.size[1]*2 if self.geom.size[1] > 0 else 100,
+            width=self.geom.size[0]*2 if self.geom.size[0] > 0 else 100, 
-                                                              depth=0.001,
+            height=self.geom.size[1]*2 if self.geom.size[1] > 0 else 100,
-                                                              create_uv_map=True,
+            depth=0.001,
-                                                              map_texture_to_each_face=True)
+            create_uv_map=True,
+            map_texture_to_each_face=True
+        )
 
         self.prim_mesh.translate(-self.prim_mesh.get_center())
 
@@ -560,17 +666,21 @@ class USDPlaneMesh(USDPrimitiveMesh):
         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 USDLight:
+class USDSphereLight:
-    def __init__(self,
+    def __init__(
-                 stage: Usd.Stage,
+        self,
-                 objid: int,
+        stage: Usd.Stage,
-                 radius: Optional[float] = 0.7):
+        objid: int,
+        radius: Optional[float] = 0.3
+    ):
         self.stage = stage
 
         xform_path = f'/World/Light_Xform_{objid}'
@@ -587,11 +697,13 @@ class USDLight:
         # defining ops required by update function
         self.translate_op = self.usd_xform.AddTranslateOp()
 
-    def update(self,
+    def update(
-               pos: np.array,
+        self,
-               intensity: int,
+        pos: np.ndarray,
-               color: np.array,
+        intensity: int,
-               frame: int):
+        color: np.ndarray,
+        frame: int
+    ):
         self.translate_op.Set(Gf.Vec3d(pos.tolist()), frame)
 
         if not np.any(pos):
@@ -600,10 +712,40 @@ class USDLight:
         self.usd_light.GetIntensityAttr().Set(intensity)
         self.usd_light.GetColorAttr().Set(Gf.Vec3d(color.tolist()))
 
+class USDDomeLight:
+    def __init__(
+        self,
+        stage: Usd.Stage,
+        objid: int
+    ):
+        self.stage = stage
+
+        xform_path = f'/World/Light_Xform_{objid}'
+        light_path = f'{xform_path}/Light_{objid}'
+        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()))
+
 class USDCamera:
-    def __init__(self,
+    def __init__(
-                 stage: Usd.Stage,
+        self,
-                 objid: int):
+        stage: Usd.Stage,
+        objid: int
+    ):
         self.stage = stage
 
         xform_path = f'/World/Camera_Xform_{objid}'
@@ -615,15 +757,20 @@ class USDCamera:
         # 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(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,
+    def update(
-               cam_pos: np.array,
+        self,
-               cam_mat: np.array,
+        cam_pos: np.ndarray,
-               frame: int):
+        cam_mat: np.ndarray,
+        frame: int
+    ):
         
         transformation_mat = 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/usd_exporter.py

@@ -0,0 +1,385 @@
+import os
+import json
+from typing import Optional, List, Union, Tuple
+
+import pprint
+from tqdm import tqdm
+from PIL import ImageOps
+from PIL import Image as im
+from pxr import Gf, Sdf, Vt
+from pxr import Usd, UsdGeom
+from termcolor import colored
+from scipy.spatial.transform import Rotation as R
+
+import mujoco
+from mujoco import mjtGeom
+from mujoco.usd.usd_utils import *
+from mujoco.usd.usd_component import *
+from mujoco import mjv_averageCamera
+from mujoco import _structs, _constants, _enums
+
+
+class USDExporter:
+
+    def __init__(
+        self,
+        model: _structs.MjModel,
+        height: int = 480,
+        width: int = 480,
+        max_geom: int = 10000,
+        output_directory_name: str = "mujoco_usdpkg",
+        output_directory_root: str = "./",
+        light_intensity: int = 10000,
+        camera_names: List[str] = None,
+        specialized_materials_file: str = None,
+        verbose: bool = True,
+    ):
+        """ Initializes a new USD Renderer
+        Args:
+            model: an MjModel instance.
+            height: image height in pixels.
+            width: image width in pixels.
+            max_geom: Optional integer specifying the maximum number of geoms that can
+                be rendered in the same scene. If None this will be chosen automatically
+                based on the estimated maximum number of renderable geoms in the model.
+            output_directory_name: name of root directory to store outputted frames and assets generated by the USD renderer.
+            output_directory_root: path to root directory storing generated frames and assets by the USD renderer.
+            verbose: decides whether to print updates.
+        """
+
+        buffer_width = model.vis.global_.offwidth
+        buffer_height = model.vis.global_.offheight
+
+        if width > buffer_width:
+            raise ValueError(f"""
+                Image width {width} > framebuffer width {buffer_width}. Either reduce the image
+                width or specify a larger offscreen framebuffer in the model XML using the
+                clause:
+                <visual>
+                <global offwidth="my_width"/>
+                </visual>""".lstrip())
+
+        if height > buffer_height:
+            raise ValueError(f"""
+                Image height {height} > framebuffer height {buffer_height}. Either reduce the
+                image height or specify a larger offscreen framebuffer in the model XML using
+                the clause:
+                <visual>
+                <global offheight="my_height"/>
+                </visual>""".lstrip())
+
+        self.model = model
+        self.height = height
+        self.width = width
+        self.max_geom = max_geom
+        self.output_directory_name = output_directory_name
+        self.output_directory_root = output_directory_root
+        self.light_intensity = light_intensity
+        self.camera_names = camera_names
+        self.specialized_materials_file = specialized_materials_file
+        self.verbose = verbose
+
+        # assert specialized_materials_file.endswith('.json')
+        # self.specialized_materials = json.loads(specialized_materials_file)
+
+        self.frame_count = 0 # maintains how many times we have saved the scene
+        self.updates = 0
+
+        self.geom_name2usd = {}
+
+        # initializing rendering requirements
+        self.renderer = mujoco.Renderer(model, height, width, max_geom)
+        self._initialize_usd_stage()
+        self._scene_option = _structs.MjvOption() # using default scene option
+
+        # initializing output_directories
+        self._initialize_output_directories()
+
+        # loading required textures for the scene
+        self._load_textures()
+
+        self.extra_added_water = False
+        self.extra_added_stove = False
+        self.extra_added_coffee = False
+
+    @property
+    def usd(self):
+        return self.stage.GetRootLayer().ExportToString()
+    
+    @property
+    def scene(self):
+        return self.renderer.scene        
+
+    def _initialize_usd_stage(self):
+        self.stage = Usd.Stage.CreateInMemory()
+        UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
+        self.stage.SetStartTimeCode(0)
+        # add as user imput
+        self.stage.SetTimeCodesPerSecond(24.0)
+
+        default_prim = UsdGeom.Xform.Define(self.stage, Sdf.Path("/World")).GetPrim()
+        self.stage.SetDefaultPrim(default_prim)
+
+    def _initialize_output_directories(self):
+        self.output_directory_path = os.path.join(self.output_directory_root, self.output_directory_name)
+        if not os.path.exists(self.output_directory_path):
+            os.makedirs(self.output_directory_path)
+
+        self.frames_directory = os.path.join(self.output_directory_path, "frames")
+        if not os.path.exists(self.frames_directory):
+            os.makedirs(self.frames_directory)
+        
+        self.assets_directory = os.path.join(self.output_directory_path, "assets")
+        if not os.path.exists(self.assets_directory):
+            os.makedirs(self.assets_directory)
+
+        if self.verbose:
+            print(colored(f"Writing output frames and assets to {self.output_directory_path}", "green"))
+
+    def update_scene(
+        self,
+        data: _structs.MjData,
+        scene_option: Optional[_structs.MjvOption] = None,
+    ):
+        """ Updates the scene with latest sim data
+        Args:
+            data: structure storing current simulation state
+            scene_option: we use this to determine which geom groups to activate
+        """
+
+        self.frame_count += 1
+
+        scene_option = scene_option or self._scene_option
+
+        # 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()
+
+            self._load_lights()
+            self._load_cameras()
+
+        self._update_geoms()
+        self._update_lights()
+        self._update_cameras(data, scene_option=scene_option)
+
+        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(range(self.model.ntex)):
+            texture_height = self.model.tex_height[texture_id]
+            texture_width = self.model.tex_width[texture_id]
+            pixels = 3*texture_height*texture_width
+            img = im.fromarray(self.model.tex_rgb[data_adr:data_adr+pixels].reshape(texture_height, texture_width, 3))
+            img = ImageOps.flip(img)
+
+            texture_file_name = f"texture_{texture_id}.png"
+
+            img.save(os.path.join(self.assets_directory, texture_file_name))
+
+            relative_path = os.path.relpath(self.assets_directory, self.frames_directory)
+            img_path = os.path.join(relative_path, texture_file_name) # relative path, TODO: switch back to this
+
+            self.texture_files.append(img_path)
+
+            data_adr += pixels
+
+        if self.verbose:
+            print(colored(f"Completed writing {self.model.ntex} textures to {self.assets_directory}", "green"))
+
+    def _load_geom(
+        self,
+        geom: _structs.MjvGeom
+    ):
+                
+        geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
+        assert geom_name not in self.geom_name2usd
+
+        # handles meshes in scene
+        if geom.type == mjtGeom.mjGEOM_MESH:
+            usd_geom = USDMesh(stage=self.stage,
+                                model=self.model,
+                                geom=geom,
+                                objid=geom_name,
+                                dataid=self.model.geom_dataid[geom.objid],
+                                rgba=geom.rgba,
+                                texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+        elif geom.type == mjtGeom.mjGEOM_PLANE:
+            usd_geom = USDPlaneMesh(stage=self.stage,
+                                    geom=geom,
+                                    objid=geom_name,
+                                    rgba=geom.rgba,
+                                    texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+        elif geom.type == mjtGeom.mjGEOM_SPHERE:
+            usd_geom = USDSphereMesh(stage=self.stage,
+                                        geom=geom,
+                                        objid=geom_name,
+                                        rgba=geom.rgba,
+                                        texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+        elif geom.type == mjtGeom.mjGEOM_CAPSULE:
+            usd_geom = USDCapsule(stage=self.stage,
+                                    geom=geom,
+                                    objid=geom_name,
+                                    rgba=geom.rgba,
+                                    texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+        elif geom.type == mjtGeom.mjGEOM_ELLIPSOID:
+            usd_geom = USDEllipsoid(stage=self.stage,
+                                    geom=geom,
+                                    objid=geom_name,
+                                    rgba=geom.rgba,
+                                    texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+        elif geom.type == mjtGeom.mjGEOM_CYLINDER:
+            usd_geom = USDCylinderMesh(stage=self.stage,
+                                        geom=geom,
+                                        objid=geom_name,
+                                        rgba=geom.rgba,
+                                        texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+        elif geom.type == mjtGeom.mjGEOM_BOX:
+            usd_geom = USDCubeMesh(stage=self.stage,
+                                    geom=geom,
+                                    objid=geom_name,
+                                    rgba=geom.rgba,
+                                    texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+        else:
+            usd_geom = None
+
+        self.geom_name2usd[geom_name] = usd_geom
+
+    def _update_geoms(self):
+
+        geom_names = set(self.geom_name2usd.keys())
+
+        # iterate through all geoms in the scene and makes update
+        for i in range(self.scene.ngeom):
+            geom = self.scene.geoms[i]
+            geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
+
+            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 self.geom_name2usd[geom_name]:
+                self.geom_name2usd[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
+        self.usd_lights = []
+        for i in range(self.scene.nlight):
+            light = self.scene.lights[i]
+            if not np.allclose(light.pos, [0, 0, 0]):
+                self.usd_lights.append(USDSphereLight(stage=self.stage,
+                                                    objid=i))
+            else:
+                self.usd_lights.append(None)
+                        
+    def _update_lights(self):
+        for i in range(self.scene.nlight):
+            light = self.scene.lights[i]
+            if not np.allclose(light.pos, [0, 0, 0]):
+                self.usd_lights[i].update(pos=light.pos,
+                                          intensity=self.light_intensity,
+                                          color=light.diffuse,
+                                          frame=self.updates)
+                
+        print("done updating")
+            
+    def _load_cameras(self):
+        self.usd_cameras = []
+        for name in self.camera_names:            
+            self.usd_cameras.append(USDCamera(stage=self.stage,
+                                    objid=name))
+        
+    def _update_cameras(
+        self,
+        data: _structs.MjData,
+        scene_option: Optional[_structs.MjvOption] = None
+    ):
+        for i in range(len(self.usd_cameras)):
+
+            camera = self.usd_cameras[i]
+            camera_name = self.camera_names[i]
+
+            self.renderer.update_scene(data, 
+                                        scene_option=scene_option,
+                                        camera=camera_name)
+            
+            avg_camera = mjv_averageCamera(self.scene.camera[0], self.scene.camera[1])
+
+            forward = avg_camera.forward
+            up = avg_camera.up
+            right = np.cross(forward, up)
+
+            R = np.eye(3)
+            R[:, 0] = right
+            R[:, 1] = up
+            R[:, 2] = -forward
+
+            camera.update(cam_pos=avg_camera.pos,
+                          cam_mat=R,
+                          frame=self.updates)
+        
+    def add_light(
+        self,
+        pos: List[float],
+        intensity:int,
+        radius: Optional[float] = 1.0,
+        color: Optional[np.array] = np.array([0.3, 0.3, 0.3]),
+        objid: Optional[int]=1,
+        light_type: Optional[str]="sphere"
+    ):
+        
+        if light_type == "sphere": 
+            new_light = USDSphereLight(stage=self.stage,
+                                       objid=objid,
+                                       radius=radius)
+
+            new_light.update(pos=pos,
+                             intensity=intensity,
+                             color=color,
+                             frame=0)
+        elif light_type == "dome":
+            new_light = USDDomeLight(stage=self.stage,
+                                     objid=objid)
+
+            new_light.update(intensity=intensity,
+                             color=color,
+                             frame=0)
+        
+    def add_camera(
+        self, 
+        pos:List[float], 
+        rotation_xyz:List[float],
+        objid: Optional[int]=1
+    ):
+        new_camera = USDCamera(stage=self.stage,
+                               objid=objid)
+        
+        r = R.from_euler('xyz', rotation_xyz, degrees=True)
+        new_camera.update(cam_pos=pos,
+                          cam_mat=r.as_matrix(),
+                          frame=0)
+        
+    def save_scene(
+        self,
+        filetype: str = "usd"
+    ):
+        self.stage.SetEndTimeCode(self.frame_count)
+        self.stage.Export(f'{self.output_directory_root}/{self.output_directory_name}/frames/frame_{self.frame_count}_.{filetype}')
+        if self.verbose:
+            print(colored(f"Writing frame_{self.frame_count}", "green"))

python/mujoco/usd_exporter.py

@@ -1,322 +0,0 @@
-import os
-import pprint
-import mujoco
-from usd_utils import *
-from PIL import ImageOps
-from mujoco import mjtGeom
-from PIL import Image as im
-from usd_component import *
-from pxr import Usd, UsdGeom
-from termcolor import colored
-from mujoco import mjv_averageCamera
-from typing import Optional, List, Union, Tuple
-from mujoco import _structs, _constants, _enums
-from scipy.spatial.transform import Rotation as R
-
-class USDExporter:
-
-    def __init__(
-        self,
-        model: _structs.MjModel,
-        height: int = 480,
-        width: int = 480,
-        max_geom: int = 10000,
-        output_directory_name: str = "mujoco_usdpkg",
-        output_directory_root: str = "./",
-        verbose: bool = True,
-        light_intensity: int = 10000
-    ):
-        """ Initializes a new USD Renderer
-        Args:
-            model: an MjModel instance.
-            height: image height in pixels.
-            width: image width in pixels.
-            max_geom: Optional integer specifying the maximum number of geoms that can
-                be rendered in the same scene. If None this will be chosen automatically
-                based on the estimated maximum number of renderable geoms in the model.
-            output_directory_name: name of root directory to store outputted frames and assets generated by the USD renderer.
-            output_directory_root: path to root directory storing generated frames and assets by the USD renderer.
-            verbose: decides whether to print updates.
-        """
-
-        buffer_width = model.vis.global_.offwidth
-        buffer_height = model.vis.global_.offheight
-
-        if width > buffer_width:
-            raise ValueError(f"""
-                Image width {width} > framebuffer width {buffer_width}. Either reduce the image
-                width or specify a larger offscreen framebuffer in the model XML using the
-                clause:
-                <visual>
-                <global offwidth="my_width"/>
-                </visual>""".lstrip())
-
-        if height > buffer_height:
-            raise ValueError(f"""
-                Image height {height} > framebuffer height {buffer_height}. Either reduce the
-                image height or specify a larger offscreen framebuffer in the model XML using
-                the clause:
-                <visual>
-                <global offheight="my_height"/>
-                </visual>""".lstrip())
-
-        self.model = model
-        self.height = height
-        self.width = width
-        self.max_geom = max_geom
-        self.output_directory_name = output_directory_name
-        self.output_directory_root = output_directory_root
-        self.verbose = verbose
-        self.light_intensity = light_intensity
-
-        self.frame_count = 0 # maintains how many times we have saved the scene
-        self.updates = 0
-
-        # initializing rendering requirements
-        self.renderer = mujoco.Renderer(model, height, width, max_geom)
-        self._initialize_usd_stage()
-        self._scene_option = _structs.MjvOption() # using default scene option
-
-        # initializing output_directories
-        self._initialize_output_directories()
-
-        # loading required textures for the scene
-        self._load_textures()
-
-    @property
-    def usd(self):
-        return self.stage.GetRootLayer().ExportToString()
-    
-    @property
-    def scene(self):
-        return self.renderer.scene        
-
-    def _initialize_usd_stage(self):
-        self.stage = Usd.Stage.CreateInMemory()
-        UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
-        self.stage.SetStartTimeCode(0)
-        # add as user imput
-        self.stage.SetTimeCodesPerSecond(60.0)
-
-    def _initialize_output_directories(self):
-        self.output_directory_path = os.path.join(self.output_directory_root, self.output_directory_name)
-        if not os.path.exists(self.output_directory_path):
-            os.makedirs(self.output_directory_path)
-
-        self.frames_directory = os.path.join(self.output_directory_path, "frames")
-        if not os.path.exists(self.frames_directory):
-            os.makedirs(self.frames_directory)
-        
-        self.assets_directory = os.path.join(self.output_directory_path, "assets")
-        if not os.path.exists(self.assets_directory):
-            os.makedirs(self.assets_directory)
-
-        if self.verbose:
-            print(colored(f"Writing output frames and assets to {self.output_directory_path}", "green"))
-
-    def update_scene(
-        self,
-        data: _structs.MjData,
-        camera: Union[int, str, _structs.MjvCamera] = -1,
-        scene_option: Optional[_structs.MjvOption] = None,
-    ):
-        """ Updates the scene with latest sim data
-        Args:
-            data: structure storing current simulation state
-            scene_option: we use this to determine which geom groups to activate
-        """
-
-        self.frame_count += 1
-
-        scene_option = scene_option or self._scene_option
-
-        # update the mujoco renderer
-        self.renderer.update_scene(data, 
-                                   scene_option=scene_option,
-                                   camera=camera)
-
-        # TODO: update scene options
-        if self.updates == 0:
-            self._initialize_usd_stage()
-
-            self._load_geoms()
-            self._load_lights()
-            self._load_cameras()
-
-        self._update_geoms()
-        self._update_lights()
-        self._update_cameras()
-
-        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 range(self.model.ntex):
-            texture_height = self.model.tex_height[texture_id]
-            texture_width = self.model.tex_width[texture_id]
-            pixels = 3*texture_height*texture_width
-            img = im.fromarray(self.model.tex_rgb[data_adr:data_adr+pixels].reshape(texture_height, texture_width, 3))
-            img = ImageOps.flip(img)
-
-            texture_file_name = f"texture_{texture_id}.png"
-
-            img.save(os.path.join(self.assets_directory, texture_file_name))
-
-            relative_path = os.path.relpath(self.assets_directory, self.frames_directory)
-            img_path = os.path.join(relative_path, texture_file_name) # relative path, TODO: switch back to this
-
-            self.texture_files.append(img_path)
-
-            data_adr += pixels
-
-            if self.verbose:
-                print(colored(f"Writing texture {texture_id}", "cyan"))
-
-        if self.verbose:
-            print(colored(f"Completed writing {self.model.ntex} textures to {self.assets_directory}", "green"))
-
-    def _load_geoms(self):
-        # stores a list of all the geoms in the scene
-        self.usd_geoms = []
-
-        # initializing the geoms
-        for i in range(self.scene.ngeom):
-            geom = self.scene.geoms[i]
-
-            if geom.rgba[3] <= 0:
-                self.usd_geoms.append(None)
-                continue
-
-            # handles meshes in scene
-            if geom.type == mjtGeom.mjGEOM_MESH:
-                usd_geom = USDMesh(stage=self.stage,
-                                   model=self.model,
-                                   geom=geom,
-                                   objid=i,
-                                   dataid=self.model.geom_dataid[geom.objid],
-                                   rgba=geom.rgba,
-                                   texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
-            # handles primitives
-            else:
-                if geom.type == mjtGeom.mjGEOM_PLANE:
-                    usd_geom = USDPlaneMesh(stage=self.stage,
-                                            geom=geom,
-                                            objid=i,
-                                            rgba=geom.rgba,
-                                            texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
-                elif geom.type == mjtGeom.mjGEOM_SPHERE:
-                    usd_geom = USDSphereMesh(stage=self.stage,
-                                             geom=geom,
-                                             objid=i,
-                                             rgba=geom.rgba,
-                                             texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
-                elif geom.type == mjtGeom.mjGEOM_CAPSULE:
-                    usd_geom = USDCapsule(stage=self.stage,
-                                          geom=geom,
-                                          objid=i,
-                                          rgba=geom.rgba,
-                                          texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
-                elif geom.type == mjtGeom.mjGEOM_ELLIPSOID:
-                    usd_geom = USDEllipsoid(stage=self.stage,
-                                            geom=geom,
-                                            objid=i,
-                                            rgba=geom.rgba,
-                                            texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
-                elif geom.type == mjtGeom.mjGEOM_CYLINDER:
-                    usd_geom = USDCylinderMesh(stage=self.stage,
-                                               geom=geom,
-                                               objid=i,
-                                               rgba=geom.rgba,
-                                               texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
-                elif geom.type == mjtGeom.mjGEOM_BOX:
-                    usd_geom = USDCubeMesh(stage=self.stage,
-                                           geom=geom,
-                                           objid=i,
-                                           rgba=geom.rgba,
-                                           texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
-                else:
-                    usd_geom = None
-            self.usd_geoms.append(usd_geom)
-
-    def _update_geoms(self):
-        # iterate through all geoms in the scene and makes update
-        for i in range(self.scene.ngeom):
-            geom = self.scene.geoms[i]
-            if self.usd_geoms[i]:
-                self.usd_geoms[i].update(pos=geom.pos,
-                                         mat=geom.mat,
-                                         frame=self.updates)
-                    
-    def _load_lights(self):
-        # initializes an usd light object for every light in the scene
-        self.usd_lights = []
-        for i in range(self.scene.nlight):
-            light = self.scene.lights[i]
-            self.usd_lights.append(USDLight(stage=self.stage,
-                                            objid=i))
-                        
-    def _update_lights(self):
-        for i in range(self.scene.nlight):
-            light = self.scene.lights[i]
-            self.usd_lights[i].update(pos=light.pos,
-                                      intensity=self.light_intensity,
-                                      color=light.diffuse,
-                                      frame=self.updates)
-            
-    def _load_cameras(self):
-        self.camera = USDCamera(stage=self.stage,
-                                objid=0)
-        
-    def _update_cameras(self):
-        camera = mjv_averageCamera(self.scene.camera[0], self.scene.camera[1])
-
-        forward = camera.forward
-        up = camera.up
-        right = np.cross(forward, up)
-
-        R = np.eye(3)
-        R[:, 0] = right
-        R[:, 1] = up
-        R[:, 2] = -forward
-
-        self.camera.update(cam_pos=camera.pos,
-                           cam_mat=R,
-                           frame=self.updates)
-        
-    def add_light(self,
-                  pos: List[float],
-                  intensity:int,
-                  radius: Optional[float] = 1.0,
-                  color: Optional[np.array] = np.array([0.3, 0.3, 0.3]),
-                  objid: Optional[int]=1):
-        new_light = USDLight(stage=self.stage,
-                             objid=objid,
-                             radius=radius)
-
-        new_light.update(pos=pos,
-                         intensity=intensity,
-                         color=color,
-                         frame=0)
-        
-    def add_camera(self, 
-                   pos:List[float], 
-                   rotation_xyz:List[float],
-                   objid: Optional[int]=1):
-        # TODO: change this!
-        new_camera = USDCamera(stage=self.stage,
-                               objid=objid)
-        
-        r = R.from_euler('xyz', rotation_xyz, degrees=True)
-        new_camera.update(cam_pos=pos,
-                          cam_mat=r.as_matrix(),
-                          frame=0)
-        
-    def save_scene(self):
-        self.stage.SetEndTimeCode(self.frame_count)
-        # with open(f'./{self.output_directory_name}/frames/frame_{self.frame_count}_.usd', "w") as f:
-        #     f.write(self.usd)
-        self.stage.Export(f'./{self.output_directory_name}/frames/frame_{self.frame_count}_.usd')
-        if self.verbose:
-            print(colored(f"Writing frame_{self.frame_count}", "green"))#