jaxsim 0.2.dev56__py3-none-any.whl → 0.2.dev77__py3-none-any.whl

jaxsim/mujoco/model.py

@@ -0,0 +1,352 @@
+import functools
+import pathlib
+from typing import Any
+
+import mujoco as mj
+import numpy as np
+import numpy.typing as npt
+from scipy.spatial.transform import Rotation
+
+
+class MujocoModelHelper:
+    """
+    Helper class to create and interact with Mujoco models and data objects.
+    """
+
+    def __init__(self, model: mj.MjModel, data: mj.MjData | None = None) -> None:
+        """"""
+
+        self.model = model
+        self.data = data if data is not None else mj.MjData(self.model)
+
+        # Populate the data with kinematics
+        mj.mj_forward(self.model, self.data)
+
+        # Keep the cache of this method local to improve GC
+        self.mask_qpos = functools.cache(self._mask_qpos)
+
+    @staticmethod
+    def build_from_xml(
+        mjcf_description: str | pathlib.Path, assets: dict[str, Any] = None
+    ) -> "MujocoModelHelper":
+        """"""
+
+        # Read the XML description if it's a path to file
+        mjcf_description = (
+            mjcf_description.read_text()
+            if isinstance(mjcf_description, pathlib.Path)
+            else mjcf_description
+        )
+
+        # Create the Mujoco model from the XML and, optionally, the assets dictionary
+        model = mj.MjModel.from_xml_string(xml=mjcf_description, assets=assets)  # noqa
+
+        return MujocoModelHelper(model=model, data=mj.MjData(model))
+
+    def time(self) -> float:
+        """Return the simulation time."""
+
+        return self.data.time
+
+    def timestep(self) -> float:
+        """Return the simulation timestep."""
+
+        return self.model.opt.timestep
+
+    def gravity(self) -> npt.NDArray:
+        """Return the 3D gravity vector."""
+
+        return self.model.opt.gravity
+
+    # =========================
+    # Methods for the base link
+    # =========================
+
+    def is_floating_base(self) -> bool:
+        """Return true if the model is floating-base."""
+
+        # A body with no joints is considered a fixed-base model.
+        # In fact, in mujoco, a floating-base model has a 6 DoFs first joint.
+        if self.number_of_joints() == 0:
+            return False
+
+        # We just check that the first joint has 6 DoFs.
+        joint0_type = self.model.jnt_type[0]
+        return joint0_type == mj.mjtJoint.mjJNT_FREE
+
+    def is_fixed_base(self) -> bool:
+        """Return true if the model is fixed-base."""
+
+        return not self.is_floating_base()
+
+    def base_link(self) -> str:
+        """Return the name of the base link."""
+
+        return mj.mj_id2name(
+            self.model, mj.mjtObj.mjOBJ_BODY, 0 if self.is_fixed_base() else 1
+        )
+
+    def base_position(self) -> npt.NDArray:
+        """Return the 3D position of the base link."""
+
+        return (
+            self.data.qpos[:3]
+            if self.is_floating_base()
+            else self.body_position(body_name=self.base_link())
+        )
+
+    def base_orientation(self, dcm: bool = False) -> npt.NDArray:
+        """Return the orientation of the base link."""
+
+        return (
+            (
+                np.reshape(self.data.xmat[0], newshape=(3, 3))
+                if dcm is True
+                else self.data.xquat[0]
+            )
+            if self.is_floating_base()
+            else self.body_orientation(body_name=self.base_link(), dcm=dcm)
+        )
+
+    def set_base_position(self, position: npt.NDArray) -> None:
+        """Set the 3D position of the base link."""
+
+        if self.is_fixed_base():
+            raise ValueError("The position of a fixed-base model cannot be set.")
+
+        position = np.atleast_1d(np.array(position).squeeze())
+
+        if position.size != 3:
+            raise ValueError(f"Wrong position size ({position.size})")
+
+        self.data.qpos[:3] = position
+
+    def set_base_orientation(self, orientation: npt.NDArray, dcm: bool = False) -> None:
+        """Set the 3D position of the base link."""
+
+        if self.is_fixed_base():
+            raise ValueError("The orientation of a fixed-base model cannot be set.")
+
+        orientation = (
+            np.atleast_2d(np.array(orientation).squeeze())
+            if dcm
+            else np.atleast_1d(np.array(orientation).squeeze())
+        )
+
+        if orientation.shape != ((4,) if not dcm else (3, 3)):
+            raise ValueError(f"Wrong orientation shape {orientation.shape}")
+
+        def is_quaternion(Q):
+            return np.allclose(np.linalg.norm(Q), 1.0)
+
+        def is_dcm(R):
+            return np.allclose(np.linalg.det(R), 1.0) and np.allclose(
+                R.T @ R, np.eye(3)
+            )
+
+        if not (is_quaternion(orientation) if not dcm else is_dcm(orientation)):
+            raise ValueError("The orientation is not a valid element of SO(3)")
+
+        W_Q_B = (
+            Rotation.from_matrix(orientation).as_quat(canonical=True)[
+                np.array([3, 0, 1, 2])
+            ]
+            if dcm
+            else orientation
+        )
+
+        self.data.qpos[3:7] = W_Q_B
+
+    # ==================
+    # Methods for joints
+    # ==================
+
+    def number_of_joints(self) -> int:
+        """"""
+
+        return self.model.njnt
+
+    def number_of_dofs(self) -> int:
+        """"""
+
+        return self.model.nq
+
+    def joint_names(self) -> list[str]:
+        """"""
+
+        return [
+            mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_JOINT, idx)
+            for idx in range(0 if self.is_fixed_base() else 1, self.number_of_joints())
+        ]
+
+    def joint_dofs(self, joint_name: str) -> int:
+        """"""
+
+        if joint_name not in self.joint_names():
+            raise ValueError(f"Joint '{joint_name}' not found")
+
+        return self.data.joint(joint_name).qpos.size
+
+    def joint_position(self, joint_name: str) -> npt.NDArray:
+        """"""
+
+        if joint_name not in self.joint_names():
+            raise ValueError(f"Joint '{joint_name}' not found")
+
+        return self.data.joint(joint_name).qpos
+
+    def joint_positions(self, joint_names: list[str] | None = None) -> npt.NDArray:
+        """"""
+
+        joint_names = joint_names if joint_names is not None else self.joint_names()
+
+        return np.hstack(
+            [self.joint_position(joint_name) for joint_name in joint_names]
+        )
+
+    def set_joint_position(
+        self, joint_name: str, position: npt.NDArray | float
+    ) -> None:
+        """"""
+
+        position = np.atleast_1d(np.array(position).squeeze())
+
+        if position.size != self.joint_dofs(joint_name=joint_name):
+            raise ValueError(
+                f"Wrong position size ({position.size}) of "
+                f"{self.joint_dofs(joint_name=joint_name)}-DoFs joint '{joint_name}'."
+            )
+
+        idx = mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_JOINT, joint_name)
+        offset = self.model.jnt_qposadr[idx]
+
+        sl = np.s_[offset : offset + self.joint_dofs(joint_name=joint_name)]
+        self.data.qpos[sl] = position
+
+    def set_joint_positions(
+        self, joint_names: list[str], positions: npt.NDArray | list[npt.NDArray]
+    ) -> None:
+        """"""
+
+        mask = self.mask_qpos(joint_names=tuple(joint_names))
+        self.data.qpos[mask] = positions
+
+    # ==================
+    # Methods for bodies
+    # ==================
+
+    def number_of_bodies(self) -> int:
+        """"""
+
+        return self.model.nbody
+
+    def body_names(self) -> list[str]:
+        """"""
+
+        return [
+            mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_BODY, idx)
+            for idx in range(self.number_of_bodies())
+        ]
+
+    def body_position(self, body_name: str) -> npt.NDArray:
+        """"""
+
+        if body_name not in self.body_names():
+            raise ValueError(f"Body '{body_name}' not found")
+
+        return self.data.body(body_name).xpos
+
+    def body_orientation(self, body_name: str, dcm: bool = False) -> npt.NDArray:
+        """"""
+
+        if body_name not in self.body_names():
+            raise ValueError(f"Body '{body_name}' not found")
+
+        return (
+            self.data.body(body_name).xmat if dcm else self.data.body(body_name).xquat
+        )
+
+    # ======================
+    # Methods for geometries
+    # ======================
+
+    def number_of_geometries(self) -> int:
+        """"""
+
+        return self.model.ngeom
+
+    def geometry_names(self) -> list[str]:
+        """"""
+
+        return [
+            mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_GEOM, idx)
+            for idx in range(self.number_of_geometries())
+        ]
+
+    def geometry_position(self, geometry_name: str) -> npt.NDArray:
+        """"""
+
+        if geometry_name not in self.geometry_names():
+            raise ValueError(f"Geometry '{geometry_name}' not found")
+
+        return self.data.geom(geometry_name).xpos
+
+    def geometry_orientation(
+        self, geometry_name: str, dcm: bool = False
+    ) -> npt.NDArray:
+        """"""
+
+        if geometry_name not in self.geometry_names():
+            raise ValueError(f"Geometry '{geometry_name}' not found")
+
+        R = np.reshape(self.data.geom(geometry_name).xmat, newshape=(3, 3))
+
+        if dcm:
+            return R
+
+        q_xyzw = Rotation.from_matrix(R).as_quat(canonical=True)
+        return q_xyzw[[3, 0, 1, 2]]
+
+    # ===============
+    # Private methods
+    # ===============
+
+    def _mask_qpos(self, joint_names: tuple[str, ...]) -> npt.NDArray:
+        """
+        Create a mask to access the DoFs of the desired `joint_names` in the `qpos` array.
+
+        Args:
+            joint_names: A tuple containing the names of the joints.
+
+        Returns:
+            A 1D array containing the indices of the `qpos` array to access the DoFs of
+            the desired `joint_names`.
+
+        Note:
+            This method takes a tuple of strings because we cache the output mask for
+            each combination of joint names. We need a hashable object for the cache.
+        """
+
+        # Get the indices of the joints in `joint_names`.
+        idxs = [
+            mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_JOINT, joint_name)
+            for joint_name in joint_names
+        ]
+
+        # We first get the index of each joint in the qpos array, and for those that
+        # have multiple DoFs, we expand their mask by appending new elements.
+        # Finally, we flatten the list of arrays to a single array, that is the
+        # final qpos mask accessing all the DoFs of the desired `joint_names`.
+        return np.atleast_1d(
+            np.hstack(
+                [
+                    np.array(
+                        [
+                            self.model.jnt_qposadr[idx] + i
+                            for i in range(self.joint_dofs(joint_name=joint_name))
+                        ]
+                    )
+                    for idx, joint_name in zip(idxs, joint_names)
+                ]
+            ).squeeze()
+        )

jaxsim/mujoco/visualizer.py

@@ -0,0 +1,152 @@
+import contextlib
+import pathlib
+from typing import ContextManager
+
+import mediapy as media
+import mujoco as mj
+import mujoco.viewer
+import numpy.typing as npt
+
+
+class MujocoVideoRecorder:
+    """"""
+
+    def __init__(
+        self,
+        model: mj.MjModel,
+        data: mj.MjData,
+        fps: int = 30,
+        width: int | None = None,
+        height: int | None = None,
+        **kwargs,
+    ) -> None:
+        """"""
+
+        width = width if width is not None else model.vis.global_.offwidth
+        height = height if height is not None else model.vis.global_.offheight
+
+        if model.vis.global_.offwidth != width:
+            model.vis.global_.offwidth = width
+
+        if model.vis.global_.offheight != height:
+            model.vis.global_.offheight = height
+
+        self.fps = fps
+        self.frames: list[npt.NDArray] = []
+        self.data: mujoco.MjData | None = None
+        self.model: mujoco.MjModel | None = None
+        self.reset(model=model, data=data)
+
+        self.renderer = mujoco.Renderer(
+            model=self.model,
+            **(dict(width=width, height=height) | kwargs),
+        )
+
+    def reset(
+        self, model: mj.MjModel | None = None, data: mj.MjData | None = None
+    ) -> None:
+        """"""
+
+        self.frames = []
+
+        self.data = data if data is not None else self.data
+        self.model = model if model is not None else self.model
+
+    def render_frame(self, camera_name: str | None = None) -> None:
+        """"""
+
+        mujoco.mj_forward(self.model, self.data)
+        self.renderer.update_scene(data=self.data)  # TODO camera name
+
+        self.frames.append(self.renderer.render())
+
+    def write_video(self, path: pathlib.Path, exist_ok: bool = False) -> None:
+        """"""
+
+        if path.is_dir():
+            raise IsADirectoryError(f"The path '{path}' is a directory.")
+
+        if not exist_ok and path.is_file():
+            raise FileExistsError(f"The file '{path}' already exists.")
+
+        media.write_video(path=path, images=self.frames, fps=self.fps)
+
+    @staticmethod
+    def compute_down_sampling(original_fps: int, target_min_fps: int) -> int:
+        """
+        Return the integer down-sampling factor to reach at least the target fps.
+
+        Args:
+            original_fps: The original fps.
+            target_min_fps: The target minimum fps.
+
+        Returns:
+            The down-sampling factor.
+        """
+
+        down_sampling = 1
+        down_sampling_final = down_sampling
+
+        while original_fps / (down_sampling + 1) >= target_min_fps:
+            down_sampling = down_sampling + 1
+
+            if int(original_fps / down_sampling) == original_fps / down_sampling:
+                down_sampling_final = down_sampling
+
+        return down_sampling_final
+
+
+class MujocoVisualizer:
+    """"""
+
+    def __init__(
+        self, model: mj.MjModel | None = None, data: mj.MjData | None = None
+    ) -> None:
+        """"""
+
+        self.data = data
+        self.model = model
+
+    def sync(
+        self,
+        viewer: mujoco.viewer.Handle,
+        model: mj.MjModel | None = None,
+        data: mj.MjData | None = None,
+    ) -> None:
+        """"""
+
+        data = data if data is not None else self.data
+        model = model if model is not None else self.model
+
+        mj.mj_forward(model, data)
+        viewer.sync()
+
+    def open_viewer(
+        self, model: mj.MjModel | None = None, data: mj.MjData | None = None
+    ) -> mj.viewer.Handle:
+        """"""
+
+        data = data if data is not None else self.data
+        model = model if model is not None else self.model
+
+        handle = mj.viewer.launch_passive(
+            model, data, show_left_ui=False, show_right_ui=False
+        )
+
+        return handle
+
+    @contextlib.contextmanager
+    def open(
+        self,
+        model: mj.MjModel | None = None,
+        data: mj.MjData | None = None,
+        close_on_exit: bool = True,
+    ) -> ContextManager[mujoco.viewer.Handle]:
+        """"""
+
+        handle = self.open_viewer(model=model, data=data)
+
+        try:
+            yield handle
+        finally:
+            handle.close() if close_on_exit else None