Merge pull request #1714 from abhihjoshi:tendons

PiperOrigin-RevId: 651048034
Change-Id: If5b7cb55ff0007c3c6f932e62f068d1168778a33

python/mujoco/usd/shapes.py

@@ -0,0 +1,191 @@
+# 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.
+# ==============================================================================
+"""Built-in shapes for USD exporter."""
+
+import mujoco
+import numpy as np
+import open3d as o3d
+
+
+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
+    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()
+  pcd.points = o3d.utility.Vector3dVector(points)
+
+  mesh = pcd.compute_convex_hull()[0]
+
+  return mesh
+
+
+def decouple_config(config: dict[str, any]):
+  """Breaks a shape config into is subcomponent shapes."""
+  decoupled_config = []
+  for key, value in config.items():
+    if key == "name":
+      continue
+    decoupled_config.append({
+        "parent_name": config["name"],
+        "name": config["name"] + "_" + key,
+        key: value.copy(),
+    })
+
+  return decoupled_config
+
+
+def mesh_config_generator(
+    name: str,
+    geom_type: mujoco.mjtGeom,
+    size: np.ndarray,
+    decouple: bool = False,
+):
+  """Creates a config for a particular mesh."""
+  if geom_type == mujoco.mjtGeom.mjGEOM_PLANE:
+    config = {
+        "name": name,
+        "box": {
+            "width": size[0] * 2 if size[0] > 0 else 100,
+            "height": size[1] * 2 if size[1] > 0 else 100,
+            "depth": 0.001,
+            "map_texture_to_each_face": True,
+        },
+    }
+  elif geom_type == mujoco.mjtGeom.mjGEOM_SPHERE:
+    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)
+    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])},
+        },
+    }
+  elif geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
+    sphere = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_SPHERE, [1.0])
+    sphere["sphere"]["transform"] = {"scale": tuple(size)}
+    config = {
+        "name": name,
+        "sphere": sphere["sphere"],
+    }
+  elif geom_type == mujoco.mjtGeom.mjGEOM_CYLINDER:
+    config = {
+        "name": name,
+        "cylinder": {
+            "radius": size[0],
+            "height": size[2] * 2,
+        },
+    }
+  elif geom_type == mujoco.mjtGeom.mjGEOM_BOX:
+    config = {
+        "name": name,
+        "box": {
+            "width": size[0] * 2,
+            "height": size[1] * 2,
+            "depth": size[2] * 2,
+        },
+    }
+  else:
+    raise NotImplementedError(
+        f"{geom_type} primitive geom type not implemented with USD integration"
+    )
+
+  if decouple:
+    config = decouple_config(config)
+
+  return config
+
+
+def mesh_generator(
+    mesh_config: dict[str, any],
+    resolution: int = 100,
+):
+  """Generates a mesh given a config consisting of shapes."""
+  assert "name" in mesh_config
+
+  prim_mesh, mesh = None, None
+
+  for shape, config in mesh_config.items():
+
+    if "name" in shape:
+      continue
+
+    if "box" in shape:
+      prim_mesh = o3d.geometry.TriangleMesh.create_box(
+          width=mesh_config[shape]["width"],
+          height=mesh_config[shape]["height"],
+          depth=mesh_config[shape]["depth"],
+          create_uv_map=True,
+          map_texture_to_each_face=True,
+      )
+    elif "hemisphere" in shape:
+      prim_mesh = create_hemisphere(radius=mesh_config[shape]["radius"])
+    elif "sphere" in shape:
+      prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
+          radius=mesh_config[shape]["radius"],
+          resolution=resolution,
+          create_uv_map=True,
+      )
+    elif "cylinder" in shape:
+      prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
+          radius=mesh_config[shape]["radius"],
+          height=mesh_config[shape]["height"],
+          resolution=resolution,
+          create_uv_map=True,
+      )
+
+    if "transform" in config:
+      if "rotate" in config["transform"]:
+        rotation = np.zeros(9)
+        quat = np.zeros(4)
+        euler = config["transform"]["rotate"]
+        seq = "xyz"
+        mujoco.mju_euler2Quat(quat, euler, seq)
+        mujoco.mju_quat2Mat(rotation, quat)
+        rotation = rotation.reshape((3, 3))
+        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"])
+        )
+      if "translate" in config["transform"]:
+        prim_mesh.translate(config["transform"]["translate"])
+
+    if not mesh:
+      mesh = prim_mesh
+    else:
+      mesh += prim_mesh
+
+  return mesh_config["name"], mesh