mink 0.0.1__py3-none-any.whl

mink/configuration.py

@@ -0,0 +1,215 @@
+from typing import Optional
+
+import mujoco
+import numpy as np
+
+from . import constants as consts
+from . import exceptions
+from .lie import SE3, SO3
+
+
+class Configuration:
+    """A struct that provides convenient access to kinematic quantities such as frame
+    transforms and frame jacobians. Frames can be defined at bodies, geoms or sites.
+
+    The `update` function ensures the proper forward kinematics functions have been
+    called, namely:
+
+    - mujoco.mj_kinematics(model, data)
+    - mujoco.mj_comPos(model, data)
+    """
+
+    def __init__(
+        self,
+        model: mujoco.MjModel,
+        q: Optional[np.ndarray] = None,
+    ):
+        """Constructor.
+
+        Args:
+            model: An instance of MjModel.
+            q: Optional configuration to initialize from. If None, the configuration
+                is initialized to the reference configuration `qpos0`.
+        """
+        self.model = model
+        self.data = mujoco.MjData(model)
+        self.update(q=q)
+
+    def update(self, q: Optional[np.ndarray] = None) -> None:
+        """Run forward kinematics.
+
+        The minimal function call required to get updated frame transforms (aka forward
+        kinematics) is `mj_kinematics`. An extra call to `mj_comPos` is needed for
+        updated Jacobians.
+
+        Args:`
+            q: Optional configuration vector to override internal data.qpos with.
+        """
+        if q is not None:
+            self.data.qpos = q
+        mujoco.mj_kinematics(self.model, self.data)
+        mujoco.mj_comPos(self.model, self.data)
+
+    def update_from_keyframe(self, key_name: str) -> None:
+        """Update the configuration from a keyframe.
+
+        Args:
+            key_name: The name of the keyframe.
+
+        Raises:
+            ValueError: if no key named `key` was found in the model.
+        """
+        key_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_KEY, key_name)
+        if key_id == -1:
+            raise exceptions.InvalidKeyframe(key_name, self.model)
+        self.update(q=self.model.key_qpos[key_id])
+
+    def check_limits(self, tol: float = 1e-6, safety_break: bool = True) -> None:
+        """Check that the current configuration is within bounds.
+
+        Args:
+            tol: Tolerance in [rad].
+            safety_break: If True, stop execution and raise an exception if the current
+                configuration is outside limits. If False, print a warning and continue
+                execution.
+        """
+        for jnt in range(self.model.njnt):
+            jnt_type = self.model.jnt_type[jnt]
+            if (
+                jnt_type == mujoco.mjtJoint.mjJNT_FREE
+                or not self.model.jnt_limited[jnt]
+            ):
+                continue
+            padr = self.model.jnt_qposadr[jnt]
+            qval = self.q[padr]
+            qmin = self.model.jnt_range[jnt, 0]
+            qmax = self.model.jnt_range[jnt, 1]
+            if qval < qmin - tol or qval > qmax + tol:
+                if safety_break:
+                    raise exceptions.NotWithinConfigurationLimits(
+                        joint_id=jnt,
+                        value=qval,
+                        lower=qmin,
+                        upper=qmax,
+                        model=self.model,
+                    )
+                else:
+                    print(
+                        f"Value {qval} at index {jnt} is out of limits: "
+                        f"[{qmin}, {qmax}]"
+                    )
+
+    def get_frame_jacobian(self, frame_name: str, frame_type: str) -> np.ndarray:
+        """Compute the Jacobian matrix of a frame velocity.
+
+        Denoting our frame by B and the world frame by W, the Jacobian matrix B_J_WB
+        is related to the body velocity B_v_WB by:
+
+            B_v_WB = B_J_WB q_dot
+
+        Args:
+            frame_name: Name of the frame in the MJCF.
+            frame_type: Type of frame. Can be a geom, a body or a site.
+
+        Returns:
+            Jacobian B_J_WB of the frame.
+        """
+        if frame_type not in consts.SUPPORTED_FRAMES:
+            raise exceptions.UnsupportedFrame(frame_type, consts.SUPPORTED_FRAMES)
+
+        frame_id = mujoco.mj_name2id(
+            self.model, consts.FRAME_TO_ENUM[frame_type], frame_name
+        )
+        if frame_id == -1:
+            raise exceptions.InvalidFrame(
+                frame_name=frame_name,
+                frame_type=frame_type,
+                model=self.model,
+            )
+
+        jac = np.empty((6, self.model.nv))
+        jac_func = consts.FRAME_TO_JAC_FUNC[frame_type]
+        jac_func(self.model, self.data, jac[:3], jac[3:], frame_id)
+
+        # MuJoCo jacobians have a frame of reference centered at the local frame but
+        # aligned with the world frame. To obtain a jacobian expressed in the local
+        # frame, aka body jacobian, we need to left-multiply by A[T_fw].
+        xmat = getattr(self.data, consts.FRAME_TO_XMAT_ATTR[frame_type])[frame_id]
+        R_wf = SO3.from_matrix(xmat.reshape(3, 3))
+        A_fw = SE3.from_rotation(R_wf.inverse()).adjoint()
+        jac = A_fw @ jac
+
+        return jac
+
+    def get_transform_frame_to_world(self, frame_name: str, frame_type: str) -> SE3:
+        """Get the pose of a frame in the current configuration.
+
+        Denoting our frame by B and the world frame by W, this function returns T_WB.
+
+        Args:
+            frame_name: Name of the frame in the MJCF.
+            frame_type: Type of frame. Can be a geom, a body or a site.
+
+        Returns:
+            The pose of the frame in the world frame.
+        """
+        if frame_type not in consts.SUPPORTED_FRAMES:
+            raise exceptions.UnsupportedFrame(frame_type, consts.SUPPORTED_FRAMES)
+
+        frame_id = mujoco.mj_name2id(
+            self.model, consts.FRAME_TO_ENUM[frame_type], frame_name
+        )
+        if frame_id == -1:
+            raise exceptions.InvalidFrame(
+                frame_name=frame_name,
+                frame_type=frame_type,
+                model=self.model,
+            )
+
+        xpos = getattr(self.data, consts.FRAME_TO_POS_ATTR[frame_type])[frame_id]
+        xmat = getattr(self.data, consts.FRAME_TO_XMAT_ATTR[frame_type])[frame_id]
+        return SE3.from_rotation_and_translation(
+            rotation=SO3.from_matrix(xmat.reshape(3, 3)),
+            translation=xpos,
+        )
+
+    def integrate(self, velocity: np.ndarray, dt: float) -> np.ndarray:
+        """Integrate a velocity starting from the current configuration.
+
+        Args:
+            velocity: The velocity.
+            dt: Integration duration in [s].
+
+        Returns:
+            The new configuration after integration.
+        """
+        q = self.data.qpos.copy()
+        mujoco.mj_integratePos(self.model, q, velocity, dt)
+        return q
+
+    def integrate_inplace(self, velocity: np.ndarray, dt: float) -> None:
+        """Integrate a velocity and update the current configuration inplace.
+
+        Args:
+            velocity: The velocity.
+            dt: Integration duration in [s].
+        """
+        mujoco.mj_integratePos(self.model, self.data.qpos, velocity, dt)
+        self.update()
+
+    # Aliases.
+
+    @property
+    def q(self) -> np.ndarray:
+        """The current configuration vector."""
+        return self.data.qpos.copy()
+
+    @property
+    def nv(self) -> int:
+        """The dimension of the tangent space."""
+        return self.model.nv
+
+    @property
+    def nq(self) -> int:
+        """The dimension of the configuration space."""
+        return self.model.nq

mink/constants.py

@@ -0,0 +1,34 @@
+import mujoco
+
+SUPPORTED_FRAMES = ("body", "geom", "site")
+
+FRAME_TO_ENUM = {
+    "body": mujoco.mjtObj.mjOBJ_BODY,
+    "geom": mujoco.mjtObj.mjOBJ_GEOM,
+    "site": mujoco.mjtObj.mjOBJ_SITE,
+}
+FRAME_TO_JAC_FUNC = {
+    "body": mujoco.mj_jacBody,
+    "geom": mujoco.mj_jacGeom,
+    "site": mujoco.mj_jacSite,
+}
+FRAME_TO_POS_ATTR = {
+    "body": "xpos",
+    "geom": "geom_xpos",
+    "site": "site_xpos",
+}
+FRAME_TO_XMAT_ATTR = {
+    "body": "xmat",
+    "geom": "geom_xmat",
+    "site": "site_xmat",
+}
+
+
+def dof_width(joint_type: mujoco.mjtJoint) -> int:
+    """Get the dimensionality of the joint in qvel."""
+    return {0: 6, 1: 3, 2: 1, 3: 1}[joint_type]
+
+
+def qpos_width(joint_type: mujoco.mjtJoint) -> int:
+    """Get the dimensionality of the joint in qpos."""
+    return {0: 7, 1: 4, 2: 1, 3: 1}[joint_type]

mink/exceptions.py

@@ -0,0 +1,98 @@
+"""Exceptions specific to mink."""
+
+from typing import Sequence
+
+import mujoco
+
+
+class MinkError(Exception):
+    """Base class for Mink exceptions."""
+
+
+class UnsupportedFrame(MinkError):
+    """Exception raised when a frame type is unsupported."""
+
+    def __init__(self, frame_type: str, supported_types: Sequence[str]):
+        message = (
+            f"{frame_type} is not supported."
+            f"Supported frame types are: {supported_types}"
+        )
+        super().__init__(message)
+
+
+class InvalidFrame(MinkError):
+    """Exception raised when a frame name is not found in the robot model."""
+
+    def __init__(
+        self,
+        frame_name: str,
+        frame_type: str,
+        model: mujoco.MjModel,
+    ):
+        if frame_type == "body":
+            available_names_of_type_frame_type = [
+                model.body(i).name for i in range(model.nbody)
+            ]
+        elif frame_type == "site":
+            available_names_of_type_frame_type = [
+                model.site(i).name for i in range(model.nsite)
+            ]
+        else:
+            assert frame_type == "geom"
+            available_names_of_type_frame_type = [
+                model.geom(i).name for i in range(model.ngeom)
+            ]
+
+        message = (
+            f"{frame_type} '{frame_name}' does not exist in the model. "
+            f"Available {frame_type} names: {available_names_of_type_frame_type}"
+        )
+
+        super().__init__(message)
+
+
+class InvalidKeyframe(MinkError):
+    """Exception raised when a keyframe name is not found in the robot model."""
+
+    def __init__(self, keyframe_name: str, model: mujoco.MjModel):
+        available_keyframes = [model.key(i).name for i in range(model.nkey)]
+        message = (
+            f"Keyframe {keyframe_name} does not exist in the model. "
+            f"Available keyframe names: {available_keyframes}"
+        )
+        super().__init__(message)
+
+
+class InvalidMocapBody(MinkError):
+    """Exception raised when a body is not a mocap body."""
+
+    def __init__(self, mocap_name: str, model: mujoco.MjModel):
+        available_mocap_names = [
+            model.body(i).name
+            for i in range(model.nbody)
+            if model.body(i).mocapid[0] != -1
+        ]
+        message = (
+            f"Body '{mocap_name}' is not a mocap body. "
+            f"Available mocap bodies: {available_mocap_names}"
+        )
+        super().__init__(message)
+
+
+class NotWithinConfigurationLimits(MinkError):
+    """Exception raised when a configuration violates its limits."""
+
+    def __init__(
+        self,
+        joint_id: int,
+        value: float,
+        lower: float,
+        upper: float,
+        model: mujoco.MjModel,
+    ):
+        joint_name = model.joint(joint_id).name
+        message = (
+            f"Joint {joint_id} ({joint_name}) violates configuration limits "
+            f"{lower} <= {value} <= {upper}"
+        )
+        super().__init__(message)

mink/lie/__init__.py

@@ -0,0 +1,13 @@
+from .base import MatrixLieGroup
+from .se3 import SE3
+from .so3 import SO3
+from .utils import get_epsilon, mat2quat, skew
+
+__all__ = (
+    "SE3",
+    "SO3",
+    "MatrixLieGroup",
+    "get_epsilon",
+    "mat2quat",
+    "skew",
+)

mink/lie/base.py

@@ -0,0 +1,156 @@
+import abc
+from typing import Union, overload
+
+import numpy as np
+from typing_extensions import Self
+
+
+class MatrixLieGroup(abc.ABC):
+    """Interface definition for matrix Lie groups.
+
+    Attributes:
+        matrix_dim: Dimension of square matrix output.
+        parameters_dim:: Dimension of underlying parameters.
+        tangent_dim:: Dimension of tangent space.
+        space_dim: Dimension of coordinates that can be transformed.
+    """
+
+    matrix_dim: int
+    parameters_dim: int
+    tangent_dim: int
+    space_dim: int
+
+    @overload
+    def __matmul__(self, other: Self) -> Self: ...
+
+    @overload
+    def __matmul__(self, other: np.ndarray) -> np.ndarray: ...
+
+    def __matmul__(self, other: Union[Self, np.ndarray]) -> Union[Self, np.ndarray]:
+        """Overload of the @ operator."""
+        if isinstance(other, np.ndarray):
+            return self.apply(target=other)
+        assert isinstance(other, MatrixLieGroup)
+        return self.multiply(other=other)
+
+    # Factory.
+
+    @classmethod
+    @abc.abstractmethod
+    def identity(cls) -> Self:
+        """Returns identity element."""
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def from_matrix(cls, matrix: np.ndarray) -> Self:
+        """Get group member from matrix representation."""
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def sample_uniform(cls) -> Self:
+        """Draw a uniform sample from the group."""
+        raise NotImplementedError
+
+    # Accessors.
+
+    @abc.abstractmethod
+    def as_matrix(self) -> np.ndarray:
+        """Get transformation as a matrix."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def parameters(self) -> np.ndarray:
+        """Get underlying representation."""
+        raise NotImplementedError
+
+    # Operations.
+
+    @abc.abstractmethod
+    def apply(self, target: np.ndarray) -> np.ndarray:
+        """Applies group action to a point."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def multiply(self, other: Self) -> Self:
+        """Composes this transformation with another."""
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def exp(cls, tangent: np.ndarray) -> Self:
+        """Computes `expm(wedge(tangent))`."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def log(self) -> np.ndarray:
+        """Computes `vee(logm(transformation matrix))`."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def adjoint(self) -> np.ndarray:
+        """Computes the adjoint."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def inverse(self) -> Self:
+        """Computes the inverse of the transform."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def normalize(self) -> Self:
+        """Normalize/projects values and returns."""
+        raise NotImplementedError
+
+    # Plus and minus operators.
+
+    # Eqn. 25.
+    def rplus(self, other: np.ndarray) -> Self:
+        return self @ self.exp(other)
+
+    # Eqn. 26.
+    def rminus(self, other: Self) -> np.ndarray:
+        return (other.inverse() @ self).log()
+
+    # Eqn. 27.
+    def lplus(self, other: np.ndarray) -> Self:
+        return self.exp(other) @ self
+
+    # Eqn. 28.
+    def lminus(self, other: Self) -> np.ndarray:
+        return (self @ other.inverse()).log()
+
+    def plus(self, other: np.ndarray) -> Self:
+        """Alias for rplus."""
+        return self.rplus(other)
+
+    def minus(self, other: Self) -> np.ndarray:
+        """Alias for rminus."""
+        return self.rminus(other)
+
+    # Jacobians.
+
+    @classmethod
+    @abc.abstractmethod
+    def ljac(cls, other: np.ndarray) -> np.ndarray:
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def ljacinv(cls, other: np.ndarray) -> np.ndarray:
+        # NOTE: Can just be np.linalg.inv(cls.ljac(other)).
+        raise NotImplementedError
+
+    # Eqn. 67.
+    @classmethod
+    def rjac(cls, other: np.ndarray) -> np.ndarray:
+        return cls.ljac(-other)
+
+    @classmethod
+    def rjacinv(cls, other: np.ndarray) -> np.ndarray:
+        return cls.ljacinv(-other)
+
+    # Eqn. 79.
+    def jlog(self) -> np.ndarray:
+        return self.rjacinv(self.log())