Add mesh loading

mojo/elements/body.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+import quaternion
 from mujoco_utils import mjcf_utils
 from typing_extensions import Self
 
@@ -77,6 +78,13 @@ def get_quaternion(self) -> np.ndarray:
             return self._mojo.physics.bind(self.mjcf.freejoint).qpos[3:].copy()
         return self._mojo.physics.bind(self.mjcf).quat.copy()
 
+    def set_euler(self, euler: np.ndarray):
+        self.set_quaternion(
+            quaternion.as_float_array(
+                quaternion.from_euler_angles(euler[0], euler[1], euler[2])
+            )
+        )
+
     def set_color(self, color: np.ndarray):
         for b in self.geoms:
             b.set_color(color)
@@ -101,7 +109,7 @@ def is_collidable(self) -> bool:
         return len(self.geoms) > 0 and self.geoms[0].is_collidable()
 
     def has_collided(self, other: Body = None):
-        if not other.is_kinematic() and not self.is_kinematic():
+        if other is not None and (not other.is_kinematic() and not self.is_kinematic()):
             warnings.warn("You are checking collisions of two non-kinematic bodies.")
         # If None, return true if there is any contact
         if other is None:

mojo/elements/geom.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -9,7 +10,7 @@
 from mojo.elements import body
 from mojo.elements.consts import GeomType, TextureMapping
 from mojo.elements.element import MujocoElement
-from mojo.elements.utils import has_collision, load_texture
+from mojo.elements.utils import has_collision, load_mesh, load_texture
 
 if TYPE_CHECKING:
     from mojo import Mojo
@@ -36,22 +37,42 @@ def create(
         quaternion: np.ndarray = None,
         color: np.ndarray = None,
         geom_type: GeomType = GeomType.BOX,
+        mesh_path: str = None,
+        mesh_scale: np.ndarray = None,
+        group: int = 1,
+        density: float = 1000,
     ) -> Self:
         position = np.array([0, 0, 0]) if position is None else position
         quaternion = np.array([1, 0, 0, 0]) if quaternion is None else quaternion
         size = np.array([0.1, 0.1, 0.1]) if size is None else size
         color = np.array([1, 1, 1, 1]) if color is None else color
         parent = body.Body.create(mojo) if parent is None else parent
+        if (
+            mesh_path
+            and geom_type != GeomType.MESH
+            or geom_type == GeomType.MESH
+            and mesh_path is None
+        ):
+            raise ValueError(
+                "To create mesh geom, 'mesh_file' must be defined "
+                "and 'geom_type' must be GeomType.MESH"
+            )
         new_geom = parent.mjcf.add(
             "geom",
             type=geom_type.value,
             pos=position,
             quat=quaternion,
             size=size,
             rgba=color,
+            group=group,
+            density=density,
         )
+        new_geom_obj = Geom(mojo, new_geom)
+        if mesh_path:
+            mesh_scale = np.array([1, 1, 1]) if mesh_scale is None else mesh_scale
+            new_geom_obj.set_mesh(mesh_path, mesh_scale)
         mojo.mark_dirty()
-        return Geom(mojo, new_geom)
+        return new_geom_obj
 
     @property
     def parent(self) -> "Body":
@@ -62,12 +83,29 @@ def parent(self) -> "Body":
 
     def set_position(self, position: np.ndarray):
         position = np.array(position)  # ensure is numpy array
+        if self.mjcf.parent.freejoint:
+            self._mojo.physics.bind(self.mjcf.parent.freejoint).qpos[:3] = position
         self._mojo.physics.bind(self.mjcf).pos = position
         self.mjcf.pos = position
 
     def get_position(self) -> np.ndarray:
+        if self.mjcf.parent.freejoint:
+            return self._mojo.physics.bind(self.mjcf.parent.freejoint).qpos[:3].copy()
         return self._mojo.physics.bind(self.mjcf).pos
 
+    def set_quaternion(self, quaternion: np.ndarray):
+        # wxyz
+        quaternion = np.array(quaternion)  # ensure is numpy array
+        if self.mjcf.parent.freejoint is not None:
+            self._mojo.physics.bind(self.mjcf.parent.freejoint).qpos[3:] = quaternion
+        self._mojo.physics.bind(self.mjcf).quat = quaternion
+        self.mjcf.quat = quaternion
+
+    def get_quaternion(self) -> np.ndarray:
+        if self.mjcf.parent.freejoint is not None:
+            return self._mojo.physics.bind(self.mjcf.parent.freejoint).qpos[3:].copy()
+        return self._mojo.physics.bind(self.mjcf).quat.copy()
+
     def set_color(self, color: np.ndarray):
         color = np.array(color)
         if len(color) == 3:
@@ -91,7 +129,8 @@ def set_texture(
         color: np.ndarray = None,
     ):
         # First check if we have loaded this texture
-        material = self._mojo.get_material(texture_path)
+        key_name = f"{texture_path}_{mapping.value}"
+        material = self._mojo.get_material(key_name)
         if material is None:
             material = load_texture(
                 self._mojo.root_element.mjcf,
@@ -105,13 +144,28 @@ def set_texture(
                 reflectance,
                 color,
             )
-            self._mojo.store_material(texture_path, material)
+            self._mojo.store_material(key_name, material)
         self.mjcf.material = material
         if self.mjcf.rgba is None:
             # Have a default white color for texture
             self.set_color(np.ones(4))
         self._mojo.mark_dirty()
 
+    def set_mesh(self, mesh_path: str, scale: np.ndarray = None):
+        scale = np.array([1, 1, 1]) if scale is None else scale
+        # First check if we have loaded this mesh
+        mesh = self._mojo.get_mesh(mesh_path)
+        if mesh is None:
+            mesh = load_mesh(self._mojo.root_element.mjcf, mesh_path, scale)
+            self._mojo.store_mesh(mesh_path, mesh)
+        self.mjcf.type = GeomType.MESH.value
+        self.mjcf.mesh = mesh.name
+        self.mjcf.contype = 0
+        self.mjcf.conaffinity = 0
+        self.mjcf.group = 1
+        self.mjcf.density = 0
+        self._mojo.mark_dirty()
+
     def set_collidable(self, value: bool):
         self.mjcf.contype = int(value)
         self.mjcf.conaffinity = int(value)
@@ -124,7 +178,17 @@ def is_collidable(self) -> bool:
             and self._mojo.physics.bind(self.mjcf).conaffinity == 1
         )
 
+    def is_kinematic(self) -> bool:
+        return self.mjcf.parent.freejoint is not None or len(self.mjcf.parent.joint) > 0
+
+    def set_kinematic(self, value: bool):
+        if value and not self.is_kinematic():
+            self.mjcf.parent.add("freejoint")
+            self._mojo.mark_dirty()
+
     def has_collided(self, other: Geom = None):
+        if other is not None and not other.is_kinematic() and not self.is_kinematic():
+            warnings.warn("You are checking collisions of two non-kinematic bodies.")
         # If None, return true if there is any contact
         if other is None:
             return len(self._mojo.physics.data.contact) > 0

mojo/elements/light.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+from mujoco_utils import mjcf_utils
 from typing_extensions import Self
 
 from mojo.elements import Body
@@ -14,6 +15,16 @@
 
 
 class Light(MujocoElement):
+    @staticmethod
+    def get(
+        mojo: Mojo,
+        name: str,
+        parent: MujocoElement = None,
+    ) -> Self:
+        root_mjcf = mojo.root_element.mjcf if parent is None else parent.mjcf
+        mjcf = mjcf_utils.safe_find(root_mjcf, "light", name)
+        return Light(mojo, mjcf)
+
     @staticmethod
     def create(
         mojo: Mojo,

mojo/elements/utils.py

@@ -46,13 +46,13 @@ def load_texture(
 ) -> mjcf.Element:
     tex_repeat = np.array([1, 1]) if tex_repeat is None else tex_repeat
     color = np.array([1, 1, 1, 1]) if color is None else color
-    uid = str(uuid.uuid4())
+    name = f"{uuid.uuid4()}_{mapping.value}"
     texture = mjcf_model.asset.add(
-        "texture", name=f"texture_{uid}", file=path, type=mapping.value
+        "texture", name=f"texture_{name}", file=path, type=mapping.value
     )
     material = mjcf_model.asset.add(
         "material",
-        name=f"material_{uid}",
+        name=f"material_{name}",
         texture=texture,
         texrepeat=tex_repeat,
         texuniform=str(tex_uniform).lower(),
@@ -63,3 +63,12 @@ def load_texture(
         rgba=color,
     )
     return material
+
+
+def load_mesh(
+    mjcf_model: mjcf.RootElement, path: str, scale: np.ndarray
+) -> mjcf.Element:
+    scale = np.array([1, 1, 1]) if scale is None else scale
+    uid = str(uuid.uuid4())
+    mesh = mjcf_model.asset.add("mesh", name=f"mesh_{uid}", file=path, scale=scale)
+    return mesh

mojo/mojo.py

@@ -14,6 +14,7 @@ def __init__(self, base_model_path: str, timestep: float = 0.01):
         model_mjcf = mjcf.from_path(base_model_path)
         self.root_element = MujocoModel(self, model_mjcf)
         self._texture_store: dict[str, mjcf.Element] = {}
+        self._mesh_store: dict[str, mjcf.Element] = {}
         self._dirty = True
         self._passive_dirty = False
         self._passive_viewer_handle = None
@@ -89,6 +90,12 @@ def get_material(self, path: str) -> mjcf.Element:
     def store_material(self, path: str, material_mjcf: mjcf.Element) -> mjcf.Element:
         self._texture_store[path] = material_mjcf
 
+    def get_mesh(self, path: str) -> mjcf.Element:
+        return self._mesh_store.get(path, None)
+
+    def store_mesh(self, path: str, mesh_mjcf: mjcf.Element) -> mjcf.Element:
+        self._mesh_store[path] = mesh_mjcf
+
     def load_model(
         self,
         path: str,

setup.py

@@ -19,7 +19,13 @@ def get_version(rel_path):
         raise RuntimeError("Unable to find version string.")
 
 
-core_requirements = ["mujoco", "numpy", "dm_control", "mujoco_utils"]
+core_requirements = [
+    "mujoco",
+    "numpy",
+    "dm_control",
+    "mujoco_utils",
+    "numpy-quaternion",
+]
 
 setuptools.setup(
     version=get_version("mojo/__init__.py"),