Fixing formatting for usd files and adding pylint configuration

python/mujoco/usd/.pylintrc

@@ -0,0 +1,8 @@
+[FORMAT]
+max-line-length=80
+
+[MESSAGES CONTROL]
+disable=bad-indentation,no-member,too-few-public-methods,too-many-arguments,too-many-instance-attributes,consider-using-enumerate,too-many-locals,undefined-loop-variable
+
+[REPORTS]
+output-format=colorized

python/mujoco/usd/camera.py

@@ -1,15 +1,27 @@
-from typing import List, Optional, Tuple
+# Copyright 2024 DeepMind Technologies Limited
-
+#
-import utils as utils_component
+# 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.
+# ==============================================================================
 import numpy as np
 
 from pxr import Gf
 from pxr import Usd
 from pxr import UsdGeom
 
-class USDCamera:
+import mujoco.usd.utils as utils_component
 
+class USDCamera:
+  """Class that handles the cameras in the USD scene"""
   def __init__(self, stage: Usd.Stage, obj_name: str):
     self.stage = stage
 
@@ -22,7 +34,6 @@ 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(12)
     self.usd_camera.CreateFocusDistanceAttr().Set(400)
 
@@ -31,8 +42,8 @@ class USDCamera:
     self.usd_camera.GetClippingRangeAttr().Set(Gf.Vec2f(1e-4, 1e6))
 
   def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
-
+    """Updates the position and orientation of the camera in the scene"""
     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)
+    self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)

python/mujoco/usd/demo.py

@@ -1,12 +1,26 @@
+# 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.
+# ==============================================================================
 import argparse
-from tqdm import tqdm
 from pathlib import Path
+from tqdm import tqdm
 
 import mujoco
 from mujoco.usd import exporter
 
 def generate_usd_trajectory(args):
-
+  """Generates a USD file from a mujoco trajectory."""
   # load a model to mujoco
   model_path = args.model_path
   m = mujoco.MjModel.from_xml_path(model_path)
@@ -19,8 +33,8 @@ def generate_usd_trajectory(args):
                              camera_names=args.camera_names)
 
   # step through the model for length steps
-  for i in tqdm(range(args.length)):
+  for _ in tqdm(range(args.length)):
-    for i in range(args.steps_per_frame):
+    for _ in range(args.steps_per_frame):
       mujoco.mj_step(m, d)
     exp.update_scene(d)
 
@@ -30,7 +44,7 @@ if __name__ == "__main__":
 
   parser = argparse.ArgumentParser()
 
-  parser.add_argument('--model_path', 
+  parser.add_argument('--model_path',
                       type=str,
                       required=True,
                       help='path to mjcf xml model')
@@ -40,20 +54,20 @@ if __name__ == "__main__":
                       default=100,
                       help='length of trajectory to render')
 
-  parser.add_argument('--output_directory_root', 
+  parser.add_argument('--output_directory_root',
                       type=str,
                       default="../usd_trajectories/",
                       help='location where to create usd files')
 
-  parser.add_argument('--camera_names', 
+  parser.add_argument('--camera_names',
                       type=str,
                       nargs='+',
                       help='cameras to include in usd')
 
-  parser.add_argument('--export_extension', 
+  parser.add_argument('--export_extension',
                       type=str,
                       default="usd",
-                      help='extension of exported file (can be usd, usda, or usdc)')
+                      help='extension of exported file (usd, usda, or usdc)')
 
   parser.add_argument('--steps_per_frame',
                       type=int,
@@ -62,7 +76,3 @@ if __name__ == "__main__":
 
   args = parser.parse_args()
   generate_usd_trajectory(args)
-
-  
-
-

python/mujoco/usd/exporter.py

@@ -13,6 +13,19 @@
 # limitations under the License.
 # ==============================================================================
 import os
+from typing import List, Optional
+
+import numpy as np
+import scipy
+import termcolor
+import tqdm
+
+from PIL import Image as im
+from PIL import ImageOps
+
+from pxr import Sdf
+from pxr import Usd
+from pxr import UsdGeom
 
 import mujoco
 
@@ -21,24 +34,9 @@ import mujoco.usd.objects as object_module
 import mujoco.usd.lights as light_module
 import mujoco.usd.camera as camera_module
 
-import numpy as np
-import scipy
-import termcolor
-import tqdm
-
-from typing import List, Optional, Tuple, Union
-from PIL import Image as im
-from PIL import ImageOps
-
-# TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
-# pytype: disable=module-attr
-from pxr import Sdf
-from pxr import Usd
-from pxr import UsdGeom
-
 
 class USDExporter:
-
+  """MuJoCo to USD exporter for porting scenes to external renderers."""
   def __init__(
       self,
       model: mujoco.MjModel,
@@ -58,7 +56,7 @@ class USDExporter:
         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
+        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.
@@ -107,6 +105,10 @@ class USDExporter:
     self.geom_names = set()
     self.geom_refs = {}
 
+    # initializing list of lights and cameras
+    self.usd_lights = []
+    self.usd_cameras = []
+
     # initializing rendering requirements
     self.renderer = mujoco.Renderer(model, height, width, max_geom)
     self._initialize_usd_stage()
@@ -120,13 +122,16 @@ class USDExporter:
 
   @property
   def usd(self):
+    """Returns the USD file as a string."""
     return self.stage.GetRootLayer().ExportToString()
 
   @property
   def scene(self):
+    """Returns the scene."""
     return self.renderer.scene
 
   def _initialize_usd_stage(self):
+    """Initializes a USD stage to represent the mujoco scene."""
     self.stage = Usd.Stage.CreateInMemory()
     UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
     self.stage.SetStartTimeCode(0)
@@ -139,6 +144,7 @@ class USDExporter:
     self.stage.SetDefaultPrim(default_prim)
 
   def _initialize_output_directories(self):
+    """Initializes output directories to store frames and assets"""
     self.output_directory_path = os.path.join(
         self.output_directory_root, self.output_directory_name
     )
@@ -167,7 +173,7 @@ class USDExporter:
       data: mujoco.MjData,
       scene_option: Optional[mujoco.MjvOption] = None,
   ):
-    """Updates the scene with latest sim data
+    """Updates the scene with latest sim data.
 
     Args:
         data: structure storing current simulation state
@@ -231,7 +237,7 @@ class USDExporter:
           )
       )
 
-  def _load_geom(self, geom: mujoco.MjvGeom, tendon: Optional[bool]=False):
+  def _load_geom(self, geom: mujoco.MjvGeom):
 
     geom_name = self._get_geom_name(geom)
 
@@ -272,7 +278,7 @@ class USDExporter:
         mesh_config = shapes_module.mesh_config_generator(
             name=geom_name,
             geom_type=geom.type,
-            size=geom.size   
+            size=geom.size
         )
         usd_geom = object_module.USDPrimitiveMesh(
             mesh_config=mesh_config,
@@ -316,12 +322,12 @@ class USDExporter:
 
   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
+        self.usd_lights.append(
-        (light_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)
 
@@ -343,7 +349,6 @@ class USDExporter:
       )
 
   def _load_cameras(self):
-    self.usd_cameras = []
     if self.camera_names is not None:
       for name in self.camera_names:
         self.usd_cameras.append(
@@ -355,7 +360,6 @@ class USDExporter:
       scene_option: Optional[mujoco.MjvOption] = None,
   ):
     for i in range(len(self.usd_cameras)):
-
       camera = self.usd_cameras[i]
       camera_name = self.camera_names[i]
 
@@ -370,12 +374,16 @@ class USDExporter:
       up = avg_camera.up
       right = np.cross(forward, up)
 
-      R = np.eye(3)
+      rotation = np.eye(3)
-      R[:, 0] = right
+      rotation[:, 0] = right
-      R[:, 1] = up
+      rotation[:, 1] = up
-      R[:, 2] = -forward
+      rotation[:, 2] = -forward
 
-      camera.update(cam_pos=avg_camera.pos, cam_mat=R, frame=self.updates)
+      camera.update(
+        cam_pos=avg_camera.pos,
+        cam_mat=rotation,
+        frame=self.updates
+      )
 
   def add_light(
       self,
@@ -386,16 +394,33 @@ class USDExporter:
       obj_name: Optional[str] = "light_1",
       light_type: Optional[str] = "sphere",
   ):
-
+    """Adds a user defined, fixed light.
+    
+    Args:
+        pos: position of the light in 3D space.
+        intensity: intensity of the light.
+        radius: radius of the light to be used by renderer.
+        color: color of the light.
+        obj_name: name associated with the light.
+        light_type: type of light (sphere or dome).
+    """
     if light_type == "sphere":
-      new_light = light_module.USDSphereLight(stage=self.stage, obj_name=str(objid), radius=radius)
+      new_light = light_module.USDSphereLight(
+        stage=self.stage,
+        obj_name=obj_name,
+        radius=radius)
 
-      new_light.update(pos=np.array(pos), intensity=intensity, color=color, frame=0)
+      new_light.update(
+        pos=np.array(pos),
+        intensity=intensity,
+        color=color,
+        frame=0)
     elif light_type == "dome":
       new_light = light_module.USDDomeLight(
-          stage=self.stage, obj_name=str(objid))
+        stage=self.stage,
-
+        obj_name=obj_name
-      new_light.update(intensity=intensity, color=color, frame=0)
+      )
+      new_light.update(intensity=intensity, color=color)
 
   def add_camera(
       self,
@@ -403,14 +428,22 @@ class USDExporter:
       rotation_xyz: List[float],
       obj_name: Optional[str] = "camera_1",
   ):
+    """Adds a user defined, fixed camera.
+    
+    Args:
+        pos: position of the camera in 3D space.
+        rotation_xyz: euler rotation of the camera.
+        obj_name: name associated with the camera.
+    """
     new_camera = camera_module.USDCamera(
-        stage=self.stage, obj_name=str(objid))
+        stage=self.stage, obj_name=obj_name)
 
     r = scipy.spatial.transform.Rotation.from_euler(
         "xyz", rotation_xyz, degrees=True)
     new_camera.update(cam_pos=np.array(pos), cam_mat=r.as_matrix(), frame=0)
 
   def save_scene(self, filetype: str = "usd"):
+    """Saves the scene to a USD file."""
     assert filetype in ["usd", "usda", "usdc"]
     self.stage.SetEndTimeCode(self.frame_count)
 
@@ -419,10 +452,14 @@ class USDExporter:
       geom_ref.update_visibility(False, geom_ref.last_visible_frame+1)
 
     self.stage.Export(
-        f"{self.output_directory_root}/{self.output_directory_name}/frames/frame_{self.frame_count}_.{filetype}"
+        f"{self.output_directory_root}/{self.output_directory_name}/" + \
+        f"frames/frame_{self.frame_count}.{filetype}"
     )
     if self.verbose:
-      print(termcolor.colored(f"Completed writing frame_{self.frame_count}.{filetype}", "green"))
+      print(termcolor.colored(
+        f"Completed writing frame_{self.frame_count}.{filetype}",
+        "green"
+      ))
 
   def _get_geom_name(self, geom):
     # adding id as part of name for USD file
@@ -431,7 +468,8 @@ class USDExporter:
       geom_name = "None"
     geom_name += f"_id{geom.objid}"
 
-    # adding additional naming information to differentiate between geoms and tendons
+    # adding additional naming information to differentiate
+    # between geoms and tendons
     if geom.objtype == mujoco.mjtObj.mjOBJ_GEOM:
       geom_name += "_geom"
     elif geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
@@ -439,9 +477,10 @@ class USDExporter:
 
     return geom_name
 
-  # for debugging purposes, prints all geoms in scene including those part of tendons
+  # for debugging purposes, prints all geoms in scene
+  # including those part of tendons
   def _print_scene_geom_info(self):
     for i in range(self.scene.ngeom):
       geom = self.scene.geoms[i]
       geom_name = self._get_geom_name(geom)
-      print(i, geom_name)
+      print(i, geom_name)

python/mujoco/usd/lights.py

@@ -1,4 +1,18 @@
-from typing import List, Optional, Tuple
+# 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 typing import Optional
 
 import numpy as np
 
@@ -8,7 +22,7 @@ from pxr import UsdGeom
 from pxr import UsdLux
 
 class USDSphereLight:
-
+  """Class that handles the sphere lights in the USD scene"""
   def __init__(
       self, stage: Usd.Stage, obj_name: str, radius: Optional[float] = 0.3
   ):
@@ -28,7 +42,14 @@ class USDSphereLight:
     # defining ops required by update function
     self.translate_op = self.usd_xform.AddTranslateOp()
 
-  def update(self, pos: np.ndarray, intensity: int, color: np.ndarray, frame: int):
+  def update(
+    self,
+    pos: np.ndarray,
+    intensity: int,
+    color: np.ndarray,
+    frame: int
+  ):
+    """Updates the attributes of a sphere light."""
     self.translate_op.Set(Gf.Vec3d(pos.tolist()), frame)
 
     if not np.any(pos):
@@ -39,7 +60,7 @@ class USDSphereLight:
 
 
 class USDDomeLight:
-
+  """Class that handles the dome lights in the USD scene"""
   def __init__(self, stage: Usd.Stage, obj_name: str):
     self.stage = stage
 
@@ -52,8 +73,8 @@ class USDDomeLight:
     # 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):
+  def update(self, intensity: int, color: np.ndarray):
+    """Updates the attributes of a dome light."""
     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

@@ -14,37 +14,34 @@
 # ==============================================================================
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from typing import List, Optional, Tuple
+from typing import Optional
-import pprint
+
+import numpy as np
+
+from pxr import Gf
+from pxr import Sdf
+from pxr import Usd
+from pxr import UsdGeom
+from pxr import UsdShade
+from pxr import Vt
 
 import mujoco
 
 import mujoco.usd.utils as utils_component
 import mujoco.usd.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
+  """ 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 
+  * `_get_uv_geometry(self)`: gets the nessecary UV information to
-     around an object in USD. Each subclass implements their own method to getting
+     wrap a texture around an object in USD. Each subclass implements
-     UV information as different objects are contructed in different ways.
+     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.
+  * `_get_mesh_geometry(self)`: gets the mesh geometry of an object
+     in the scene.
   """
 
   def __init__(
@@ -68,16 +65,16 @@ class USDObject(ABC):
     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
+    self.last_visible_frame = -2
 
   @abstractmethod
   def _get_uv_geometry(self):
-    """Gets UV information for an object in the scene"""
+    """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"""
+    """Gets structure of an object in the scene."""
     raise NotImplementedError
 
   def attach_image_material(self, usd_mesh):
@@ -176,13 +173,22 @@ class USDObject(ABC):
   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):
+  def update(
-    """Updates the position and orientation of an object in the scene for a given frame"""
+    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))
@@ -204,7 +210,7 @@ class USDObject(ABC):
     self.scale_op.Set(Gf.Vec3f(scale.tolist()), frame)
 
 class USDMesh(USDObject):
-
+  """Class that handles predefined meshes in the USD scene."""
   def __init__(
       self,
       stage: Usd.Stage,
@@ -249,7 +255,7 @@ class USDMesh(USDObject):
     else:
       self.attach_solid_material(self.usd_mesh)
 
-  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):
@@ -268,7 +274,7 @@ class USDMesh(USDObject):
         mesh_texcoord_adr_from:mesh_texcoord_adr_to
     ]
 
-    mesh_facetexcoord_ranges = self.get_facetexcoord_ranges(
+    mesh_facetexcoord_ranges = self._get_facetexcoord_ranges(
         self.model.nmesh, self.model.mesh_facenum
     )
 
@@ -302,7 +308,7 @@ class USDMesh(USDObject):
     return mesh_vert, mesh_face, mesh_facenum
 
 class USDPrimitiveMesh(USDObject):
-
+  """Class to handle primitive shapes in the USD scene."""
   def __init__(
       self,
       mesh_config: dict,
@@ -333,13 +339,13 @@ class USDPrimitiveMesh(USDObject):
     self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
 
     # setting mesh uv properties
-    mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
+    mesh_texcoord, _ = 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.texcoords.SetIndices(Vt.IntArray(list(range(mesh_facenum * 3))))
 
     self._set_refinement_properties(self.usd_prim)
 
@@ -349,12 +355,12 @@ class USDPrimitiveMesh(USDObject):
       self.attach_solid_material(self.usd_mesh)
 
   def generate_primitive_mesh(self):
+    """Generates the mesh for the primitive USD object."""
     _, 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
@@ -364,10 +370,12 @@ class USDPrimitiveMesh(USDObject):
     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"]
+        x_multiplier = self.mesh_config["box"]["width"]
+        y_multiplier = 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"]
+        x_multiplier = self.mesh_config["sphere"]["radius"]
-      
+        y_multiplier = self.mesh_config["sphere"]["radius"]
+
     mesh_texcoord[:, 0] *= x_scale * x_multiplier
     mesh_texcoord[:, 1] *= y_scale * y_multiplier
 
@@ -383,7 +391,7 @@ class USDPrimitiveMesh(USDObject):
     return mesh_vert, mesh_face, len(mesh_face)
 
 class USDTendon(USDObject):
-  
+  """Class to handle tendons in the USD scene."""
   def __init__(
       self,
       mesh_config: dict,
@@ -403,53 +411,70 @@ class USDTendon(USDObject):
     self.tendon_parts = self.generate_primitive_mesh()
     self.usd_refs = defaultdict(dict)
 
-    for name, mesh in self.tendon_parts.items():
+    for name, _ 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)
+      usd_xform = UsdGeom.Xform.Define(
-      self.usd_refs[name]["usd_mesh"] = UsdGeom.Mesh.Define(stage, mesh_path)
+        stage,
+        part_xform_path
+      )
+      self.usd_refs[name]["usd_xform"] = usd_xform
+      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]["translate_op"] = usd_xform.AddTranslateOp()
-      self.usd_refs[name]["scale_op"] = self.usd_refs[name]["usd_xform"].AddScaleOp()
+      self.usd_refs[name]["scale_op"] = 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"].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"])
+      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(
+      self.texcoords = UsdGeom.PrimvarsAPI(
-          "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
+        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)]))
+      self.texcoords.SetIndices(
-    
+        Vt.IntArray(list(range(part_geometry['mesh_facenum'] * 3)))
-    for name in self.usd_refs.keys(): 
+      )
-      self._set_refinement_properties(self.usd_refs[name]["usd_prim"])
+
+    for _, ref in self.usd_refs.items():
+      self._set_refinement_properties(ref["usd_prim"])
       if self.texture_file:
-        self.attach_image_material(self.usd_refs[name]["usd_mesh"])
+        self.attach_image_material(ref["usd_mesh"])
       else:
-        self.attach_solid_material(self.usd_refs[name]["usd_mesh"])
+        self.attach_solid_material(ref["usd_mesh"])
-    
+
   def generate_primitive_mesh(self):
+    """Generates the tendon mesh using primitives."""
     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] = {
@@ -471,8 +496,14 @@ class USDTendon(USDObject):
       }
     return part_geometries
 
-  def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int, scale: Optional[np.ndarray] = None):
+  def update(
-    """Updates the position and orientation of an object in the scene for a given frame"""
+    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."""
     super().update(pos, mat, visible, frame, scale)
     for name in self.tendon_parts.keys():
       if "left" in name:
@@ -483,7 +514,7 @@ class USDTendon(USDObject):
         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"""
+    """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)
@@ -491,5 +522,3 @@ class USDTendon(USDObject):
         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

@@ -1,15 +1,28 @@
-import copy
+# Copyright 2024 DeepMind Technologies Limited
-import mujoco
+#
-import pprint
+# 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.
+# ==============================================================================
 import numpy as np
 import open3d as o3d
 
+import mujoco
+
 def create_hemisphere(
     radius: float,
     theta_steps: int = 50,
     phi_steps: int = 50
 ):
-
+    """Creates a hemisphere mesh from a point cloud."""
     points = []
     for i in range(phi_steps + 1):
         phi = np.pi / 2 * i / phi_steps
@@ -28,6 +41,7 @@ def create_hemisphere(
     return mesh
 
 def decouple_config(config: dict):
+  """Breaks a shape config into is subcomponent shapes."""
   decoupled_config = []
   for key, value in config.items():
     if key == "name":
@@ -46,7 +60,7 @@ def mesh_config_generator(
     size: np.ndarray,
     decouple: bool = False
 ):
-
+    """Creates a config for a particular mesh."""
     if geom_type == mujoco.mjtGeom.mjGEOM_PLANE:
       config = {
         "name": name,
@@ -65,7 +79,11 @@ def mesh_config_generator(
         }
       }
     elif geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE:
-      cylinder = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_CYLINDER, size)
+      cylinder = mesh_config_generator(
+        name,
+        mujoco.mjtGeom.mjGEOM_CYLINDER,
+        size
+      )
       config = {
         "name": name,
         "cylinder": cylinder["cylinder"],
@@ -84,7 +102,11 @@ def mesh_config_generator(
         },
       }
     elif geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
-      sphere = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_SPHERE, [1.0])
+      sphere = mesh_config_generator(
+        name,
+        mujoco.mjtGeom.mjGEOM_SPHERE,
+        [1.0]
+      )
       sphere["sphere"]["transform"] = {
         "scale": tuple(size)
       }
@@ -110,7 +132,9 @@ def mesh_config_generator(
         }
       }
     else:
-      raise NotImplemented(f"{geom_type} primitive geom type not implemented with USD integration")
+      raise NotImplementedError(
+        f"{geom_type} primitive geom type not implemented with USD integration"
+      )
 
     if decouple:
       config = decouple_config(config)
@@ -121,10 +145,10 @@ def mesh_generator(
     mesh_config: dict,
     resolution: int = 100,
 ):
-
+    """Generates a mesh given a config consisting of shapes."""
     assert "name" in mesh_config
 
-    mesh = None
+    prim_mesh, mesh = None, None
 
     for shape, config in mesh_config.items():
 
@@ -139,7 +163,7 @@ def mesh_generator(
             create_uv_map=True,
             map_texture_to_each_face=True,
         )
-      elif "hemisphere" in shape: 
+      elif "hemisphere" in shape:
         prim_mesh = create_hemisphere(
           radius=mesh_config[shape]["radius"]
         )
@@ -159,17 +183,19 @@ def mesh_generator(
 
       if "transform" in config:
         if "rotate" in config["transform"]:
-          R = np.zeros(9)
+          rotation = np.zeros(9)
           quat = np.zeros(4)
           euler = config["transform"]["rotate"]
           seq = 'xyz'
           mujoco.mju_euler2Quat(quat, euler, seq)
-          mujoco.mju_quat2Mat(R, quat)
+          mujoco.mju_quat2Mat(rotation, quat)
-          R = R.reshape((3,3))
+          rotation = rotation.reshape((3,3))
-          prim_mesh.rotate(R, center=(0, 0, 0))
+          prim_mesh.rotate(rotation, center=(0, 0, 0))
         if "scale" in config["transform"]:
           prim_mesh.vertices = o3d.utility.Vector3dVector(
-          np.asarray(prim_mesh.vertices) * np.array(config["transform"]["scale"]))
+            np.asarray(prim_mesh.vertices) * \
+            np.array(config["transform"]["scale"])
+          )
         if "translate" in config["transform"]:
           prim_mesh.translate(config["transform"]["translate"])
 
@@ -179,10 +205,3 @@ def mesh_generator(
         mesh += prim_mesh
 
     return mesh_config["name"], mesh
-      
-
-
-
-
-    
-