38 trimesh constructors

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mujoco-mojo"
-version = "2.2.0"
+version = "2.2.1"
 description = "A complete MJCF lifecycle and trial orchestration suite for MuJoCo, powered by Pydantic v2."
 readme = "README.md"
 requires-python = ">=3.12"

src/mujoco_mojo/mjcf/mujoco_attr/asset_attr/mesh.py

@@ -1,17 +1,32 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Annotated, Any, ClassVar, Literal, Self, overload
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    ClassVar,
+    Literal,
+    NotRequired,
+    Self,
+    TypedDict,
+    Unpack,
+    overload,
+)
 
 import mujoco
 import numpy as np
+import trimesh
+import trimesh.visual
 from pydantic import Field, field_validator, model_validator
 
 from mujoco_mojo.mjcf.dependency_path import DepPath
 from mujoco_mojo.mjcf.xml_model import XMLModel
 from mujoco_mojo.typing import Inertia, MaterialName, MeshName, Vec3, Vec4
 from mujoco_mojo.utils.log import get_logger
 
+if TYPE_CHECKING:
+    from trimesh.typed import BooleanEngineType
 logger = get_logger(__name__)
 
 __all__ = [
@@ -29,21 +44,39 @@
     "name",
     "class_",
     "content_type",
+    "scale",
     "file",
     "vertex",
     "normal",
     "texcoord",
     "face",
     "refpos",
     "refquat",
-    "scale",
     "smoothnormal",
     "maxhullvert",
     "inertia",
     "material",
 )
 
 
+class MeshAttributes(TypedDict):
+    name: NotRequired[MeshName | None]
+    class_: NotRequired[str | None]
+    content_type: NotRequired[str | None]
+    file: NotRequired[DepPath | None]
+    scale: NotRequired[Vec3]
+    inertia: NotRequired[Inertia]
+    smoothnormal: NotRequired[bool]
+    maxhullvert: NotRequired[int]
+    vertex: NotRequired[tuple[tuple[float, float, float], ...] | None]
+    normal: NotRequired[tuple[tuple[float, float, float], ...] | None]
+    texcoord: NotRequired[tuple[tuple[float, float], ...] | None]
+    face: NotRequired[tuple[tuple[int, int, int], ...] | None]
+    refpos: NotRequired[Vec3]
+    refquat: NotRequired[Vec4]
+    material: NotRequired[MaterialName | None]
+
+
 class MeshBase(XMLModel):
     """Base class from which other mesh classes are built from."""
 
@@ -314,6 +347,157 @@ class Mesh(MeshBase):
 
     builtin: Literal["none"] = "none"
 
+    @classmethod
+    def new_from_trimesh(
+        cls,
+        mesh: trimesh.Trimesh | trimesh.Scene,
+        **kwargs: Unpack[MeshAttributes],
+    ) -> Self:
+        """
+        Creates a Mesh instance from a trimesh object.
+
+        Automatically extracts vertices, faces, and (if present) normals/texcoords.
+        """
+        if isinstance(mesh, trimesh.Scene):
+            mesh = mesh.to_mesh()
+
+        assert isinstance(mesh, trimesh.Trimesh)
+        mesh.process(validate=True)
+        mesh.fill_holes()
+
+        # mandatory geometry data
+        extracted_data: dict[str, Any] = {
+            "vertex": tuple(map(tuple, mesh.vertices.tolist())),
+            "face": tuple(map(tuple, mesh.faces.tolist())),
+        }
+
+        # extract normals (if they exist and weren't provided)
+        if "normal" not in kwargs:
+            # only pull if the count matches
+            v_normals = getattr(mesh, "vertex_normals", None)
+            if v_normals is not None and len(v_normals) == len(mesh.vertices):
+                extracted_data["normal"] = tuple(map(tuple, v_normals.tolist()))
+
+        # extract texture coordinates (if they exist and weren't provided)
+        if "texcoord" not in kwargs:
+            if isinstance(mesh.visual, trimesh.visual.TextureVisuals):
+                uvs = mesh.visual.uv
+                if uvs is not None and len(uvs) == len(mesh.vertices):
+                    extracted_data["texcoord"] = tuple(map(tuple, uvs.tolist()))
+
+        # merge kwargs
+        final_attrs = {**extracted_data, **kwargs}
+
+        return cls(**final_attrs)
+
+    @classmethod
+    def new_frustum(
+        cls,
+        radius_bottom: float,
+        radius_top: float,
+        height: float,
+        sections: int = 32,
+        wall_thickness: float = 0.0,
+        engine: BooleanEngineType = None,
+        base_at_zero: bool = True,
+        **kwargs: Unpack[MeshAttributes],
+    ) -> Self:
+        """
+        Generates a conical frustum (tapered cylinder).
+
+        Args:
+            radius_bottom (float): Radius at z = -height/2.
+            radius_top (float): Radius at z = +height/2.
+            height (float): Total height of the frustum.
+            sections (int): Number of radial segments.
+            wall_thickness (float): Thickness of the walls (makes a hollow frustum if not equal to zero).
+            engine (BooleanEngineType): Which engine to use for boolean subtraction (only used for hollow volumes, reccomended to use "manifold" but requires `manifold3d`).
+            base_at_zero (bool): Whether or not to shift the volume so it sits on z=0 instead of underground.
+            **kwargs (MeshAttributes): Standard MuJoCo mesh attributes (name, inertia, etc.)
+
+        Example:
+            <img src="https://raw.githubusercontent.com/Hydrowelder/mujoco-mojo/refs/heads/master/docs/assets/mesh/frustums.jpg" width="300" />
+
+            Two frustums generated using this method. The red frustum has its wall thickness set to 0 and its bottom radius is greater than the top. The cyan frustum has a non-zero wall thickness and a larger top radius than the bottom resulting in a hollow volume. The cyan frustum has more sections than the red.
+
+        """
+        # create the outer mesh
+        mesh = trimesh.creation.cylinder(radius=1.0, height=height, sections=sections)
+
+        def taper_mesh(m: trimesh.Trimesh, r_b: float, r_t: float) -> None:
+            verts = m.vertices.copy()
+            # Dynamically find the midpoint of the specific mesh's height extent
+            z_ext = m.bounds[:, 2]
+            z_mid = (z_ext[1] - z_ext[0]) / 4
+
+            mask_top = verts[:, 2] > z_mid
+            mask_bottom = verts[:, 2] < -z_mid
+
+            verts[mask_top, :2] *= r_t
+            verts[mask_bottom, :2] *= r_b
+            m.vertices = verts
+
+        taper_mesh(mesh, radius_bottom, radius_top)
+
+        if wall_thickness != 0:
+            inner = trimesh.creation.cylinder(
+                radius=1.0, height=height + 0.02, sections=sections
+            )
+            taper_mesh(
+                inner, radius_bottom - wall_thickness, radius_top - wall_thickness
+            )
+
+            # perform boolean subtraction to make it hollow
+            # Note: This requires the 'manifold' or 'blender' engine installed for trimesh,
+            # otherwise it falls back to a simpler but less robust internal boolean.
+            mesh = mesh.difference(inner, engine=engine)
+
+        if base_at_zero:
+            mesh.apply_translation([0, 0, height / 2])
+        return cls.new_from_trimesh(mesh, **kwargs)
+
+    @classmethod
+    def new_hollow_box(
+        cls,
+        width: float,
+        depth: float,
+        height: float,
+        wall_thickness: float = 0.02,
+        engine: BooleanEngineType = None,
+        base_at_zero: bool = True,
+        **kwargs: Unpack[MeshAttributes],
+    ) -> Self:
+        """
+        Generates an open-top box with physical wall thickness (concave bin).
+
+        Args:
+            width (float): Total outer dimension along X.
+            depth (float): Total outer dimension along Y.
+            height (float): Total outer dimension along Z.
+            wall_thickness (float): Thickness of the walls and bottom floor.
+            engine (BooleanEngineType): Engine to use for boolean subtraction.
+            base_at_zero (bool): Shift volume so it sits on z=0 instead of underground.
+            **kwargs (MeshAttributes): Standard MuJoCo mesh attributes (name, inertia, etc.)
+
+        """
+        # outer box centered at (0, 0, 0)
+        outer = trimesh.creation.box(extents=(width, depth, height))
+
+        # inner cutout box
+        inner = trimesh.creation.box(
+            extents=(width - (2 * wall_thickness), depth - (2 * wall_thickness), height)
+        )
+
+        # shift the inner tool up by the floor thickness
+        inner.apply_translation([0, 0, wall_thickness])
+
+        mesh = outer.difference(inner, engine=engine)
+
+        if base_at_zero:
+            mesh.apply_translation([0, 0, height / 2])
+
+        return cls.new_from_trimesh(mesh, **kwargs)
+
 
 class MeshSphere(MeshBase):
     """

src/mujoco_mojo/mjcf/mujoco_attr/body_attr/geom.py

@@ -4,6 +4,7 @@
 
 import mujoco
 import numpy as np
+import trimesh
 from pydantic import ConfigDict, Field
 
 from mujoco_mojo.mjcf.defaults import SOLIMP_DEFAULT, SOLREF_DEFAULT
@@ -262,10 +263,7 @@ def request(
         signal_manager: SignalManager,
         attrs: list[
             Literal["xpos", "xmat", "xvelp", "xvelr", "xaccp", "xaccr", "quat"]
-        ] = [
-            "xpos",
-            "quat",
-        ],
+        ] = ["xpos", "quat"],
     ):
         """
         Registers specific geom attributes for logging.
@@ -476,6 +474,31 @@ class GeomMesh(GeomBase, ProximityMixin):
     mesh: MeshName
     """If the geom type is "mesh", this attribute is required. It references the mesh asset to be instantiated. This attribute can also be specified if the geom type corresponds to a geometric primitive, namely one of "sphere", "capsule", "cylinder", "ellipsoid", "box". In that case the primitive is automatically fitted to the mesh asset referenced here. The fitting procedure uses either the equivalent inertia box or the axis-aligned bounding box of the mesh, as determined by the attribute fitaabb of compiler. The resulting size of the fitted geom is usually what one would expect, but if not, it can be further adjusted with the fitscale attribute below. In the compiled mjModel the geom is represented as a regular geom of the specified primitive type, and there is no reference to the mesh used for fitting."""
 
+    def trimesh(self, mj_model: mujoco.MjModel) -> trimesh.Trimesh:
+        # get mesh id and data from mujoco
+        mesh_id = mj_model.geom_dataid[self.get_id(mj_model)]
+
+        if mesh_id == -1:
+            msg = "Exact proximity mesh tool is not currently supported for geoms of this type. Please use the `SPHERE_TO_SPHERE`/`CONVEX_HULL` algorithm, or convert the Geom to a GeomMesh."
+            logger.error(msg)
+            raise TypeError(msg)
+
+        # extract vertices and faces
+        if self._local_verts is None:
+            adr = mj_model.mesh_vertadr[mesh_id]
+            num = mj_model.mesh_vertnum[mesh_id]
+            self._local_verts = mj_model.mesh_vert[adr : adr + num].copy()
+
+        f_adr = mj_model.mesh_faceadr[mesh_id]
+        f_num = mj_model.mesh_facenum[mesh_id]
+        self._local_faces = mj_model.mesh_face[f_adr : f_adr + f_num].copy()
+
+        # create a trimesh and its proximity query
+        self._baked_mesh = trimesh.Trimesh(
+            vertices=self._local_verts, faces=self._local_faces
+        )
+        return self._baked_mesh
+
 
 class GeomSDF(GeomBase):
     attributes = (*_geom_attr, "type")

src/mujoco_mojo/utils/proximity_mixin.py

@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
 import mujoco
@@ -15,7 +16,7 @@
 logger = get_logger(__name__)
 
 
-class ProximityMixin(MojoBaseModel):
+class ProximityMixin(MojoBaseModel, ABC):
     _baked_mesh: trimesh.Trimesh | None = PrivateAttr(default=None)
     """Internal trimesh representation of the geometry."""
 
@@ -63,6 +64,10 @@ def vertex_max_norm(self, mj_model: mujoco.MjModel) -> tuple[float, Vec3]:
 
         return radius, centroid
 
+    @abstractmethod
+    def trimesh(self, mj_model: mujoco.MjModel) -> trimesh.Trimesh:
+        """Generates a trimesh object to be used with proximity calculations."""
+
     def bake_proximity(self, mj_model: mujoco.MjModel, proximity_type: ProximityType):
         """Builds the BVH tree from the comiled MuJoCo mesh data."""
         geom_self: Proximityable = self  # pyright: ignore[reportAssignmentType]
@@ -75,28 +80,13 @@ def bake_proximity(self, mj_model: mujoco.MjModel, proximity_type: ProximityType
             case ProximityType.SPHERE_TO_SPHERE | ProximityType.CONVEX_HULL:
                 return
 
-        # get mesh id and data from mujoco
-        mesh_id = mj_model.geom_dataid[geom_self.get_id(mj_model)]
-
-        if mesh_id == -1:
-            msg = "Exact proximity mesh tool is not currently supported for geoms of this type. Please use the `SPHERE_TO_SPHERE`/`CONVEX_HULL` algorithm, or convert the Geom to a GeomMesh."
-            logger.error(msg)
-            raise TypeError(msg)
-
-        # extract vertices and faces
-        if self._local_verts is None:
-            adr = mj_model.mesh_vertadr[mesh_id]
-            num = mj_model.mesh_vertnum[mesh_id]
-            self._local_verts = mj_model.mesh_vert[adr : adr + num].copy()
-
-        f_adr = mj_model.mesh_faceadr[mesh_id]
-        f_num = mj_model.mesh_facenum[mesh_id]
-        self._local_faces = mj_model.mesh_face[f_adr : f_adr + f_num].copy()
-
-        # create a trimesh and its proximity query
-        self._baked_mesh = trimesh.Trimesh(
-            vertices=self._local_verts, faces=self._local_faces
+        self.trimesh(mj_model)
+        assert (
+            self._baked_mesh
+            and self._local_faces is not None
+            and self._local_verts is not None
         )
+
         match proximity_type:
             case ProximityType.FACE_TO_FACE:
                 if not geom_self.name: