diff --git a/docs/assets/mesh/frustums.jpg b/docs/assets/mesh/frustums.jpg new file mode 100644 index 00000000..9e832794 Binary files /dev/null and b/docs/assets/mesh/frustums.jpg differ diff --git a/pyproject.toml b/pyproject.toml index 0420e3d5..5df945b2 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/mujoco_mojo/mjcf/mujoco_attr/asset_attr/mesh.py b/src/mujoco_mojo/mjcf/mujoco_attr/asset_attr/mesh.py index c8349143..c1366c19 100644 --- a/src/mujoco_mojo/mjcf/mujoco_attr/asset_attr/mesh.py +++ b/src/mujoco_mojo/mjcf/mujoco_attr/asset_attr/mesh.py @@ -1,10 +1,23 @@ 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 @@ -12,6 +25,8 @@ 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,6 +44,7 @@ "name", "class_", "content_type", + "scale", "file", "vertex", "normal", @@ -36,7 +52,6 @@ "face", "refpos", "refquat", - "scale", "smoothnormal", "maxhullvert", "inertia", @@ -44,6 +59,24 @@ ) +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: + + + 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): """ diff --git a/src/mujoco_mojo/mjcf/mujoco_attr/body_attr/geom.py b/src/mujoco_mojo/mjcf/mujoco_attr/body_attr/geom.py index ae4c771d..5d397dd2 100644 --- a/src/mujoco_mojo/mjcf/mujoco_attr/body_attr/geom.py +++ b/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") diff --git a/src/mujoco_mojo/utils/proximity_mixin.py b/src/mujoco_mojo/utils/proximity_mixin.py index 4a952327..690e822a 100644 --- a/src/mujoco_mojo/utils/proximity_mixin.py +++ b/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: diff --git a/tests/mjcf/mesh.py b/tests/mjcf/mesh.py new file mode 100644 index 00000000..17f209fd --- /dev/null +++ b/tests/mjcf/mesh.py @@ -0,0 +1,98 @@ +from pathlib import Path + +import numpy as np + +import mujoco_mojo as mojo + +logger = mojo.utils.get_logger(__name__) + + +def generate(mojo_model: mojo.MojoModel, *args, **kwargs) -> mojo.MojoModel: + """Generates two boxes at varying distances.""" + mojo_model.mjcf.assets = [ + mojo.Asset( + meshes=[ + frustum_mesh := mojo.Mesh.new_frustum( + radius_bottom=3, + radius_top=1, + height=2, + sections=10, + name=mojo.MeshName("frustum_mesh"), + ), + hollow_frustum_mesh := mojo.Mesh.new_frustum( + radius_bottom=1, + radius_top=2, + height=2, + wall_thickness=0.1, + sections=20, + engine="manifold", + name=mojo.MeshName("hollow_frustum_mesh"), + ), + ], + ), + ] + assert frustum_mesh.name and hollow_frustum_mesh.name + + mojo_model.mjcf.worldbody = mojo.WorldBody() + + # create two boxes with sites for the springs + mojo_model.mjcf.worldbody.bodies.extend( + [ + mojo.Body( + name=mojo.BodyName("body1"), + geoms=[ + mojo.GeomMesh( + name=mojo.GeomName("frustum"), + mesh=frustum_mesh.name, + rgba=mojo.utils.Color.ROSE_500.with_alpha(0.5), + ), + ], + ), + mojo.Body( + name=mojo.BodyName("body2"), + freejoints=[mojo.FreeJoint()], + geoms=[ + mojo.GeomMesh( + name=mojo.GeomName("hollow_frustum"), + mesh=hollow_frustum_mesh.name, + rgba=mojo.utils.Color.CYAN_500.with_alpha(0.5), + pose=mojo.PoseQuat(pos=np.array([4.5, 0, 2])), + ), + ], + ), + ] + ) + + return mojo_model + + +def profile(): + import cProfile + import pstats + + mojo.utils.setup_logger() + + workdir = Path(__file__).parent / "mc_bumper_test_2" + # Run the Monte Carlo + runner = mojo.utils.MojoRunner( + generator=generate, + workdir=workdir, + config=mojo.utils.MonteCarloConfig(n_trial=1, n_proc=1, resume=False), + ) + + profiler = cProfile.Profile() + profiler.enable() + had_fails = runner.run(clean_workdir=True, cleanup_delay=-1) + profiler.disable() + + stats = pstats.Stats(profiler).sort_stats("cumulative") + stats.dump_stats(save_as := "baseline.prof") + stats.print_callees(30, "apply_load") + + print(f"To view results run:\n\tsnakeviz {save_as}") + print(f"Finished with {had_fails=}") + + +if __name__ == "__main__": + # main() + profile()