Merge pull request #1827 from abhihjoshi:open3d-removal

PiperOrigin-RevId: 662510404
Change-Id: I79b356eecddb3402c444cb16a305516bd724aef6

python/mujoco/usd/camera.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 """Camera handling for USD exporter."""
 
-import mujoco.usd.utils as utils_component
+import mujoco.usd.utils as utils_module
 
 import numpy as np
 
@@ -49,7 +49,7 @@ class USDCamera:
 
   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(
+    transformation_mat = utils_module.create_transform_matrix(
         rotation_matrix=cam_mat, translation_vector=cam_pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)

python/mujoco/usd/demo.py

@@ -42,6 +42,10 @@ def generate_usd_trajectory(local_args):
     if exp.frame_count < d.time * local_args.framerate:
       exp.update_scene(data=d)
 
+  exp.add_light(pos=(0, 0, 0),
+                intensity=2000,
+                light_type='dome')
+
   exp.save_scene(filetype=local_args.export_extension)
 
 

python/mujoco/usd/exporter.py

@@ -27,7 +27,6 @@ from PIL import Image as im
 from PIL import ImageOps
 import scipy
 import termcolor
-import tqdm
 
 # TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
 # pytype: disable=module-attr
@@ -209,7 +208,7 @@ class USDExporter:
     """Load textures."""
     data_adr = 0
     self.texture_files = []
-    for texture_id in tqdm.tqdm(range(self.model.ntex)):
+    for texture_id in range(self.model.ntex):
       texture_height = self.model.tex_height[texture_id]
       texture_width = self.model.tex_width[texture_id]
       texture_nchannel = self.model.tex_nchannel[texture_id]
@@ -251,13 +250,13 @@ class USDExporter:
 
     assert geom_name not in self.geom_names
 
-    texture_file = (
+    if geom.matid == -1:
-        self.texture_files[
+      geom_textures = []
-            self.model.mat_texid[geom.matid][mujoco.mjTEXROLE_RGB]
+    else:
-        ]
+      geom_textures = [
-        if geom.matid != -1
+          (self.texture_files[i], self.model.tex_type[i]) if i != -1 else None
-        else None
+          for i in self.model.mat_texid[geom.matid]
-    )
+      ]
 
     # handling meshes in our scene
     if geom.type == mujoco.mjtGeom.mjGEOM_MESH:
@@ -268,7 +267,7 @@ class USDExporter:
           obj_name=geom_name,
           dataid=self.model.geom_dataid[geom.objid],
           rgba=geom.rgba,
-          texture_file=texture_file,
+          geom_textures=geom_textures,
       )
     else:
       # handling tendons in our scene
@@ -282,10 +281,11 @@ class USDExporter:
         usd_geom = object_module.USDTendon(
             mesh_config=mesh_config,
             stage=self.stage,
+            model=self.model,
             geom=geom,
             obj_name=geom_name,
             rgba=geom.rgba,
-            texture_file=texture_file,
+            geom_textures=geom_textures,
         )
       # handling primitives in our scene
       else:
@@ -297,10 +297,11 @@ class USDExporter:
         usd_geom = object_module.USDPrimitiveMesh(
             mesh_config=mesh_config,
             stage=self.stage,
+            model=self.model,
             geom=geom,
             obj_name=geom_name,
             rgba=geom.rgba,
-            texture_file=texture_file,
+            geom_textures=geom_textures,
         )
 
     self.geom_names.add(geom_name)

python/mujoco/usd/objects.py

@@ -16,13 +16,14 @@
 
 import abc
 import collections
-from typing import Optional, Dict, Any
+from typing import Any, Dict, List, Optional, Tuple
 
 import mujoco
-import mujoco.usd.shapes as shapes_component
+import mujoco.usd.shapes as shapes_module
-import mujoco.usd.utils as utils_component
+import mujoco.usd.utils as utils_module
 import numpy as np
 
+
 # TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
 # pytype: disable=module-attr
 from pxr import Gf
@@ -50,16 +51,18 @@ class USDObject(abc.ABC):
   def __init__(
       self,
       stage: Usd.Stage,
+      model: mujoco.MjModel,
       geom: mujoco.MjvGeom,
       obj_name: str,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      texture_file: Optional[str] = None,
+      geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
   ):
     self.stage = stage
+    self.model = model
     self.geom = geom
     self.obj_name = obj_name
     self.rgba = rgba
-    self.texture_file = texture_file
+    self.geom_textures = geom_textures
 
     self.xform_path = f"/World/Mesh_Xform_{obj_name}"
     self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path)
@@ -117,7 +120,7 @@ class USDObject(abc.ABC):
     # setting the image texture attributes
     image_shader.CreateIdAttr("UsdUVTexture")
     image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(
-        self.texture_file
+        self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB][0]
     )
     image_shader.CreateInput("sourceColorSpace", Sdf.ValueTypeNames.Token).Set(
         "sRGB"
@@ -185,7 +188,7 @@ class USDObject(abc.ABC):
       scale: Optional[np.ndarray] = None,
   ):
     """Updates the position and orientation of an object."""
-    transformation_mat = utils_component.create_transform_matrix(
+    transformation_mat = utils_module.create_transform_matrix(
         rotation_matrix=mat, translation_vector=pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
@@ -222,11 +225,10 @@ class USDMesh(USDObject):
       obj_name: str,
       dataid: int,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      texture_file: Optional[str] = None,
+      geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
   ):
-    super().__init__(stage, geom, obj_name, rgba, texture_file)
+    super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
 
-    self.model = model
     self.dataid = dataid
 
     mesh_path = f"{self.xform_path}/Mesh_{obj_name}"
@@ -241,15 +243,19 @@ class USDMesh(USDObject):
     )
     self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
 
-    # setting mesh uv properties
+    if (
-    mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
+        geom.matid != -1
-    self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
+        and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]
-        "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
+    ):
-    )
+      # setting mesh uv properties
-    self.texcoords.Set(mesh_texcoord)
+      mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
-    self.texcoords.SetIndices(Vt.IntArray(mesh_facetexcoord.tolist()))
+      self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
-
+          "UVMap",
-    if self.texture_file:
+          Sdf.ValueTypeNames.TexCoord2fArray,
+          UsdGeom.Tokens.faceVarying,
+      )
+      self.texcoords.Set(mesh_texcoord)
+      self.texcoords.SetIndices(Vt.IntArray(mesh_facetexcoord.tolist()))
       self.attach_image_material(self.usd_mesh)
     else:
       self.attach_solid_material(self.usd_mesh)
@@ -314,12 +320,13 @@ class USDPrimitiveMesh(USDObject):
       self,
       mesh_config: Dict[Any, Any],
       stage: Usd.Stage,
+      model: mujoco.MjModel,
       geom: mujoco.MjvGeom,
       obj_name: str,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      texture_file: Optional[str] = None,
+      geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
   ):
-    super().__init__(stage, geom, obj_name, rgba, texture_file)
+    super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
 
     self.mesh_config = mesh_config
     self.prim_mesh = self.generate_primitive_mesh()
@@ -334,42 +341,66 @@ class USDPrimitiveMesh(USDObject):
         [3 for _ in range(mesh_facenum)]
     )
     self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
-
-    # setting mesh uv properties
-    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(list(range(mesh_facenum * 3))))
-
     self._set_refinement_properties(self.usd_prim)
 
-    if self.texture_file:
+    if (
+        geom.matid != -1
+        and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]
+    ):
+      # setting mesh uv properties
+      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(list(range(mesh_facenum * 3))))
+
       self.attach_image_material(self.usd_mesh)
     else:
       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)
+    tex_role = mujoco.mjtTextureRole
+    geom_rgb_texture = (
+        self.geom_textures[tex_role.mjTEXROLE_RGB]
+        if self.geom_textures
+        else None
+    )
+    texture_type = geom_rgb_texture[1] if geom_rgb_texture else None
+    _, prim_mesh = shapes_module.mesh_factory(self.mesh_config, texture_type)
     prim_mesh.translate(-prim_mesh.get_center())
     return prim_mesh
 
   def _get_uv_geometry(self):
-    assert self.prim_mesh
+    assert self.prim_mesh and self.prim_mesh.triangle_uvs is not None
 
     mesh_texcoord = np.array(self.prim_mesh.triangle_uvs)
     mesh_facetexcoord = np.asarray(self.prim_mesh.triangles)
-    # TODO(etom): bring back support for rescaling the texture coordinates.
+    tex_role = mujoco.mjtTextureRole
+    geom_rgb_texture = self.geom_textures[tex_role.mjTEXROLE_RGB][1]
+
+    if geom_rgb_texture == mujoco.mjtTexture.mjTEXTURE_2D:
+      s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid]
+
+      if self.model.mat_texuniform[self.geom.matid]:
+        if self.geom.size[0] > 0:
+          s_scale *= self.geom.size[0]
+        if self.geom.size[1] > 0:
+          t_scale *= self.geom.size[1]
+
+      mesh_texcoord[:, 0] *= s_scale / (self.geom.size[0] * 2)
+      mesh_texcoord[:, 1] *= t_scale / (self.geom.size[1] * 2)
 
     return mesh_texcoord, mesh_facetexcoord.flatten()
 
   def _get_mesh_geometry(self):
     assert self.prim_mesh
 
-    # get mesh geometry from the open3d mesh model
+    # get mesh geometry
     mesh_vert = np.asarray(self.prim_mesh.vertices)
     mesh_face = np.asarray(self.prim_mesh.triangles)
 
@@ -383,12 +414,13 @@ class USDTendon(USDObject):
       self,
       mesh_config: Dict[Any, Any],
       stage: Usd.Stage,
+      model: mujoco.MjModel,
       geom: mujoco.MjvGeom,
       obj_name: str,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      texture_file: Optional[str] = None,
+      geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
   ):
-    super().__init__(stage, geom, obj_name, rgba, texture_file)
+    super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
 
     self.mesh_config = mesh_config
     self.tendon_parts = self.generate_primitive_mesh()
@@ -407,6 +439,7 @@ class USDTendon(USDObject):
 
     # setting mesh geometry properties for each of the parts in the tendon
     part_geometries = self._get_mesh_geometry()
+    part_geometry = None
     for name, part_geometry in part_geometries.items():
       self.usd_refs[name]["usd_mesh"].GetPointsAttr().Set(
           part_geometry["mesh_vert"]
@@ -418,33 +451,43 @@ class USDTendon(USDObject):
           part_geometry["mesh_face"]
       )
 
-    # setting uv properties for each of the parts in the tendon
+    tex_role = mujoco.mjtTextureRole
-    part_uv_geometries = self._get_uv_geometry()
+    if geom.matid != -1 and self.geom_textures[tex_role.mjTEXROLE_RGB]:
-    for name, part_uv_geometry in part_uv_geometries.items():
+      # setting uv properties for each of the parts in the tendon
-      self.texcoords = UsdGeom.PrimvarsAPI(
+      part_uv_geometries = self._get_uv_geometry()
-          self.usd_refs[name]["usd_mesh"]
+      for name, part_uv_geometry in part_uv_geometries.items():
-      ).CreatePrimvar(
+        self.texcoords = UsdGeom.PrimvarsAPI(
-          "UVMap",
+            self.usd_refs[name]["usd_mesh"]
-          Sdf.ValueTypeNames.TexCoord2fArray,
+        ).CreatePrimvar(
-          UsdGeom.Tokens.faceVarying,
+            "UVMap",
-      )
+            Sdf.ValueTypeNames.TexCoord2fArray,
-      self.texcoords.Set(part_uv_geometry["mesh_texcoord"])
+            UsdGeom.Tokens.faceVarying,
-      self.texcoords.SetIndices(
+        )
-          Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3)))
+        self.texcoords.Set(part_uv_geometry["mesh_texcoord"])
-      )
+        self.texcoords.SetIndices(
-
+            Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3)))
-    for _, ref in self.usd_refs.items():
+        )
-      self._set_refinement_properties(ref["usd_prim"])
+        for _, ref in self.usd_refs.items():
-      if self.texture_file:
+          self._set_refinement_properties(ref["usd_prim"])
-        self.attach_image_material(ref["usd_mesh"])
+          self.attach_image_material(ref["usd_mesh"])
-      else:
+    else:
+      for _, ref in self.usd_refs.items():
+        self._set_refinement_properties(ref["usd_prim"])
         self.attach_solid_material(ref["usd_mesh"])
 
   def generate_primitive_mesh(self):
     """Generates the tendon mesh using primitives."""
     mesh_parts = {}
+    geom_rgb_texture = (
+        self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]
+        if self.geom_textures
+        else None
+    )
+    texture_type = geom_rgb_texture[1] if geom_rgb_texture else None
     for part_config in self.mesh_config:
-      mesh_name, prim_mesh = shapes_component.mesh_generator(part_config)
+      mesh_name, prim_mesh = shapes_module.mesh_factory(
+          part_config, texture_type
+      )
       prim_mesh.translate(-prim_mesh.get_center())
       mesh_parts[mesh_name] = prim_mesh
     return mesh_parts
@@ -452,6 +495,7 @@ class USDTendon(USDObject):
   def _get_uv_geometry(self):
     part_uv_geometries = collections.defaultdict(dict)
     for name, mesh in self.tendon_parts.items():
+      assert mesh.triangle_uvs is not None
       mesh_texcoord = np.array(mesh.triangle_uvs)
       mesh_facetexcoord = np.asarray(mesh.triangles)
       part_uv_geometries[name] = {
@@ -463,7 +507,7 @@ class USDTendon(USDObject):
   def _get_mesh_geometry(self):
     part_geometries = collections.defaultdict(dict)
     for name, mesh in self.tendon_parts.items():
-      # get mesh geometry from the open3d mesh model
+      # get mesh geometry
       mesh_vert = np.asarray(mesh.vertices)
       mesh_face = np.asarray(mesh.triangles)
       part_geometries[name] = {

python/mujoco/usd/shapes.py

@@ -14,33 +14,270 @@
 # ==============================================================================
 """Built-in shapes for USD exporter."""
 
-from typing import Dict, Any
+from typing import Any, Dict, Optional, Tuple, Union
 
 import mujoco
 import numpy as np
-from open3d import open3d as o3d
 
 
-def create_hemisphere(
+def get_triangle_uvs(
-    radius: float, theta_steps: int = 50, phi_steps: int = 50
+    vertices: np.ndarray,
+    triangles: np.ndarray,
+    texture_type: Optional[mujoco.mjtTexture]
 ):
-  """Creates a hemisphere mesh from a point cloud."""
+  """Returns UV coordinates for a given mesh."""
-  points = []
+  if not texture_type:
-  for i in range(phi_steps + 1):
+    return None
-    phi = np.pi / 2 * i / phi_steps
-    for j in range(theta_steps + 1):
-      theta = 2 * np.pi * j / theta_steps
-      x = radius * np.sin(phi) * np.cos(theta)
-      y = radius * np.sin(phi) * np.sin(theta)
-      z = radius * np.cos(phi)
-      points.append([x, y, z])
 
-  pcd = o3d.geometry.PointCloud()
+  triangle_uvs = []
-  pcd.points = o3d.utility.Vector3dVector(points)
+  if texture_type == mujoco.mjtTexture.mjTEXTURE_2D:
+    triangle_uvs = [
+        [vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)
+    ]
 
-  mesh = pcd.compute_convex_hull()[0]
+  elif texture_type == mujoco.mjtTexture.mjTEXTURE_CUBE:
+    center = np.mean(vertices, axis=0)
+    for vertex_id in np.nditer(triangles):
+      x, y, z = vertices[vertex_id] - center
 
-  return mesh
+      abs_x, abs_y, abs_z = abs(x), abs(y), abs(z)
+      u = 0
+      v = 0
+
+      if x > 0 and abs_x >= abs_y and abs_x >= abs_z:
+        u = -z / abs_x
+        v = y / abs_x
+      elif x <= 0 and abs_x >= abs_y and abs_x >= abs_z:
+        u = z / abs_x
+        v = y / abs_x
+      elif y > 0 and abs_y >= abs_x and abs_y >= abs_z:
+        u = x / abs_y
+        v = -z / abs_y
+      elif y <= 0 and abs_y >= abs_x and abs_y >= abs_z:
+        u = x / abs_y
+        v = z / abs_y
+      elif z > 0 and abs_z >= abs_x and abs_z >= abs_y:
+        u = x / abs_z
+        v = y / abs_z
+      elif z <= 0 and abs_z >= abs_x and abs_z >= abs_y:
+        u = -x / abs_z
+        v = y / abs_z
+
+      u = (u + 1.0) / 2.0
+      v = (v + 1.0) / 2.0
+      v /= 6
+
+      assert 0 <= u and u <= 1 and 0 <= v and v <= 1
+
+      triangle_uvs.append([u, v])
+  elif texture_type == mujoco.mjtTexture.mjTEXTURE_SKYBOX:
+    # defaults to 2D mapping temporarily
+    triangle_uvs = [
+        [vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)
+    ]
+
+  return np.array(triangle_uvs)
+
+
+class TriangleMesh:
+  """Store UV and geometry information for a primitive mesh.
+
+  Attributes:
+    vertices: A numpy array of vertices.
+    triangles: A numpy array of triangles.
+    triangle_uvs: A numpy array of UV coordinates.
+  """
+
+  def __init__(self,
+               vertices: np.ndarray,
+               triangles: np.ndarray,
+               triangle_uvs: np.ndarray):
+    """Creates a TriangleMesh object.
+
+    Args:
+      vertices: A numpy array of vertices.
+      triangles: A numpy array of triangles.
+      triangle_uvs: A numpy array of UV coordinates.
+    """
+    self.vertices = vertices
+    self.triangles = triangles
+    self.triangle_uvs = triangle_uvs
+
+  @classmethod
+  def create_box(
+      cls,
+      width: float,
+      height: float,
+      depth: float,
+      texture_type: Optional[mujoco.mjtTexture]
+  ) -> TriangleMesh:
+    """Creates a box."""
+    vertices = np.array([[0.0, 0.0, 0.0],
+                         [width, 0.0, 0.0],
+                         [0.0, 0.0, depth],
+                         [width, 0.0, depth],
+                         [0.0, height, 0.0],
+                         [width, height, 0.0],
+                         [0.0, height, depth],
+                         [width, height, depth]])
+
+    triangles = np.array([[4, 7, 5],
+                          [4, 6, 7],
+                          [0, 2, 4],
+                          [2, 6, 4],
+                          [0, 1, 2],
+                          [1, 3, 2],
+                          [1, 5, 7],
+                          [1, 7, 3],
+                          [2, 3, 7],
+                          [2, 7, 6],
+                          [0, 4, 1],
+                          [1, 4, 5]])
+
+    triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
+
+    return TriangleMesh(vertices, triangles, triangle_uvs)
+
+  @classmethod
+  def create_sphere(
+      cls,
+      radius: float,
+      texture_type: Optional[mujoco.mjtTexture],
+      resolution: int
+  ) -> TriangleMesh:
+    """Creates a sphere."""
+    vertices = []
+    triangles = []
+    for i in range(2*resolution + 1):
+      phi = np.pi * i / (2*resolution)
+      for j in range(resolution + 1):
+        theta = 2 * np.pi * j / resolution
+        x = radius * np.sin(phi) * np.cos(theta)
+        y = radius * np.sin(phi) * np.sin(theta)
+        z = radius * np.cos(phi)
+        vertices.append([x, y, z])
+
+    for i in range(2*resolution):
+      for j in range(resolution):
+        first = i * (resolution + 1) + j
+        second = first + resolution + 1
+
+        triangles.append([first, second, first + 1])
+        triangles.append([second, second + 1, first + 1])
+
+    vertices = np.array(vertices)
+    triangles = np.array(triangles)
+
+    triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
+
+    return TriangleMesh(vertices, triangles, triangle_uvs)
+
+  @classmethod
+  def create_hemisphere(
+      cls,
+      radius: float,
+      texture_type: Optional[mujoco.mjtTexture],
+      resolution: int,
+  ) -> TriangleMesh:
+    """Creates a hemisphere."""
+    vertices = []
+    triangles = []
+    for i in range(resolution + 1):
+      phi = np.pi / 2 * i / (resolution)
+      for j in range(resolution + 1):
+        theta = 2 * np.pi * j / resolution
+        x = radius * np.sin(phi) * np.cos(theta)
+        y = radius * np.sin(phi) * np.sin(theta)
+        z = radius * np.cos(phi)
+        vertices.append([x, y, z])
+    vertices.append([0, 0, 0])
+
+    for i in range(resolution):
+      for j in range(resolution):
+        first = i * (resolution + 1) + j
+        second = first + resolution + 1
+
+        triangles.append([first, second, first + 1])
+        triangles.append([second, second + 1, first + 1])
+
+    for i in range(resolution):
+      first = resolution * (resolution + 1) + i
+      triangles.append([first, first + 1, len(vertices) - 1])
+
+    vertices = np.array(vertices)
+    triangles = np.array(triangles)
+
+    triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
+
+    return TriangleMesh(vertices, triangles, triangle_uvs)
+
+  @classmethod
+  def create_cylinder(
+      cls,
+      radius: float,
+      height: float,
+      texture_type: Optional[mujoco.mjtTexture],
+      resolution: int
+  ) -> TriangleMesh:
+    """Creates a cylinder."""
+    vertices = []
+    triangles = []
+
+    # adding all the vertices for the cylinder including
+    # two center vertices at ends
+    for i in range(2):
+      z = 0 if i == 0 else height
+      for j in range(resolution + 1):
+        theta = 2 * np.pi * j / resolution
+        x = radius * np.cos(theta)
+        y = radius * np.sin(theta)
+        vertices.append([x, y, z])
+    vertices.append([0, 0, 0])
+    vertices.append([0, 0, height])
+
+    # constructing the end faces for the cylinder
+    for i in range(2):
+      for j in range(resolution):
+        first = (resolution + 1) * i + j
+        triangles.append([first, first + 1, len(vertices) - (2 - i)])
+
+    # constructing side of cylinder
+    for i in range(resolution):
+      second = resolution + 1 + i
+      triangles.append([i, second, second + 1])
+      triangles.append([i, i + 1, second + 1])
+
+    vertices = np.array(vertices)
+    triangles = np.array(triangles)
+
+    triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
+
+    return TriangleMesh(vertices, triangles, triangle_uvs)
+
+  def translate(self, translation: np.array):
+    self.vertices = self.vertices + translation
+
+  def rotate(self, rotation: np.array, center: Tuple[float, ...]):
+    translated_point = self.vertices - center
+    self.vertices = np.dot(translated_point, rotation) + center
+
+  def scale(self, scale: np.array):
+    self.vertices = self.vertices * scale
+
+  def get_center(self):
+    center = np.mean(self.vertices, axis=0)
+    return center
+
+  def __add__(self, other):
+    if isinstance(other, TriangleMesh):
+      new_vertices = np.vstack((self.vertices, other.vertices))
+      other_triangles = other.triangles + len(self.vertices)
+      new_triangles = np.vstack((self.triangles, other_triangles))
+      new_triangle_uvs = None
+      if self.triangle_uvs is not None:
+        new_triangle_uvs = np.vstack((self.triangle_uvs, other.triangle_uvs))
+      return TriangleMesh(new_vertices, new_triangles, new_triangle_uvs)
+    raise TypeError(f"Cannot add TriangleMesh with {type(other)}")
 
 
 def decouple_config(config: Dict[str, Any]):
@@ -60,7 +297,7 @@ def decouple_config(config: Dict[str, Any]):
 
 def mesh_config_generator(
     name: str,
-    geom_type: int | mujoco.mjtGeom,
+    geom_type: Union[int, mujoco.mjtGeom],
     size: np.ndarray,
     decouple: bool = False,
 ):
@@ -79,20 +316,22 @@ def mesh_config_generator(
     config = {"name": name, "sphere": {"radius": float(size[0])}}
   elif geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE:
     cylinder = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_CYLINDER, size)
+    cylinder["cylinder"]["transform"] = {
+        "transform": {"translate": (0, 0, size[2])}
+    }
     config = {
         "name": name,
         "cylinder": cylinder["cylinder"],
         "left_hemisphere": {
             "radius": size[0],
             "transform": {
-                "translate": (0, 0, -size[2]),
                 "rotate": (np.pi, 0, 0),
             },
         },
         "right_hemisphere": {
             "radius": size[0],
-            "transform": {"translate": (0, 0, size[2])},
+            "transform": {"translate": (0, 0, 2*size[2])},
-        },
+        }
     }
   elif geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
     sphere = mesh_config_generator(
@@ -131,8 +370,9 @@ def mesh_config_generator(
   return config
 
 
-def mesh_generator(
+def mesh_factory(
     mesh_config: Dict[str, Any],
+    texture_type: Optional[mujoco.mjtTexture],
     resolution: int = 100,
 ):
   """Generates a mesh given a config consisting of shapes."""
@@ -146,27 +386,30 @@ def mesh_generator(
       continue
 
     if "box" in shape:
-      prim_mesh = o3d.geometry.TriangleMesh.create_box(
+      prim_mesh = TriangleMesh.create_box(
           width=mesh_config[shape]["width"],
           height=mesh_config[shape]["height"],
           depth=mesh_config[shape]["depth"],
-          create_uv_map=True,
+          texture_type=texture_type
-          map_texture_to_each_face=True,
       )
     elif "hemisphere" in shape:
-      prim_mesh = create_hemisphere(radius=mesh_config[shape]["radius"])
+      prim_mesh = TriangleMesh.create_hemisphere(
-    elif "sphere" in shape:
-      prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
           radius=mesh_config[shape]["radius"],
-          resolution=resolution,
+          texture_type=texture_type,
-          create_uv_map=True,
+          resolution=resolution
+      )
+    elif "sphere" in shape:
+      prim_mesh = TriangleMesh.create_sphere(
+          radius=mesh_config[shape]["radius"],
+          texture_type=texture_type,
+          resolution=resolution
       )
     elif "cylinder" in shape:
-      prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
+      prim_mesh = TriangleMesh.create_cylinder(
           radius=mesh_config[shape]["radius"],
           height=mesh_config[shape]["height"],
-          resolution=resolution,
+          texture_type=texture_type,
-          create_uv_map=True,
+          resolution=resolution
       )
     else:
       raise ValueError("Shape not supported")
@@ -182,10 +425,7 @@ def mesh_generator(
         rotation = rotation.reshape((3, 3))
         prim_mesh.rotate(rotation, center=(0, 0, 0))
       if "scale" in config["transform"]:
-        prim_mesh.vertices = o3d.utility.Vector3dVector(
+        prim_mesh.scale(config["transform"]["scale"])
-            np.asarray(prim_mesh.vertices)
-            * np.array(config["transform"]["scale"])
-        )
       if "translate" in config["transform"]:
         prim_mesh.translate(config["transform"]["translate"])