mink 0.0.1__py3-none-any.whl

mink/lie/se3.py

@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import mujoco
+import numpy as np
+
+from ..exceptions import InvalidMocapBody
+from .base import MatrixLieGroup
+from .so3 import SO3
+from .utils import get_epsilon, skew
+
+_IDENTITY_WXYZ_XYZ = np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)
+
+
+@dataclass(frozen=True)
+class SE3(MatrixLieGroup):
+    """Special Euclidean group for proper rigid transforms in 3D.
+
+    Internal parameterization is (qw, qx, qy, qz, x, y, z). Tangent parameterization is
+    (vx, vy, vz, omega_x, omega_y, omega_z).
+    """
+
+    wxyz_xyz: np.ndarray
+    matrix_dim: int = 4
+    parameters_dim: int = 7
+    tangent_dim: int = 6
+    space_dim: int = 3
+
+    def __repr__(self) -> str:
+        quat = np.round(self.wxyz_xyz[:4], 5)
+        xyz = np.round(self.wxyz_xyz[4:], 5)
+        return f"{self.__class__.__name__}(wxyz={quat}, xyz={xyz})"
+
+    def copy(self) -> SE3:
+        return SE3(wxyz_xyz=np.array(self.wxyz_xyz))
+
+    def parameters(self) -> np.ndarray:
+        return self.wxyz_xyz
+
+    @classmethod
+    def identity(cls) -> SE3:
+        return SE3(wxyz_xyz=_IDENTITY_WXYZ_XYZ)
+
+    @classmethod
+    def from_rotation_and_translation(
+        cls,
+        rotation: SO3,
+        translation: np.ndarray,
+    ) -> SE3:
+        assert translation.shape == (SE3.space_dim,)
+        return SE3(wxyz_xyz=np.concatenate([rotation.wxyz, translation]))
+
+    @classmethod
+    def from_rotation(cls, rotation: SO3) -> SE3:
+        return SE3.from_rotation_and_translation(
+            rotation=rotation,
+            translation=np.zeros(
+                SE3.space_dim,
+            ),
+        )
+
+    @classmethod
+    def from_translation(cls, translation: np.ndarray) -> SE3:
+        return SE3.from_rotation_and_translation(
+            rotation=SO3.identity(), translation=translation
+        )
+
+    @classmethod
+    def from_matrix(cls, matrix: np.ndarray) -> SE3:
+        assert matrix.shape == (SE3.matrix_dim, SE3.matrix_dim)
+        return SE3.from_rotation_and_translation(
+            rotation=SO3.from_matrix(matrix[:3, :3]),
+            translation=matrix[:3, 3],
+        )
+
+    @classmethod
+    def from_mocap_id(cls, data: mujoco.MjData, mocap_id: int) -> SE3:
+        return SE3.from_rotation_and_translation(
+            rotation=SO3(data.mocap_quat[mocap_id]),
+            translation=data.mocap_pos[mocap_id],
+        )
+
+    @classmethod
+    def from_mocap_name(
+        cls, model: mujoco.MjModel, data: mujoco.MjData, mocap_name: str
+    ) -> SE3:
+        mocap_id = model.body(mocap_name).mocapid[0]
+        if mocap_id == -1:
+            raise InvalidMocapBody(mocap_name, model)
+        return SE3.from_mocap_id(data, mocap_id)
+
+    @classmethod
+    def sample_uniform(cls) -> SE3:
+        return SE3.from_rotation_and_translation(
+            rotation=SO3.sample_uniform(),
+            translation=np.random.uniform(-1.0, 1.0, size=(SE3.space_dim,)),
+        )
+
+    def rotation(self) -> SO3:
+        return SO3(wxyz=self.wxyz_xyz[:4])
+
+    def translation(self) -> np.ndarray:
+        return self.wxyz_xyz[4:]
+
+    def as_matrix(self) -> np.ndarray:
+        hmat = np.eye(self.matrix_dim, dtype=np.float64)
+        hmat[:3, :3] = self.rotation().as_matrix()
+        hmat[:3, 3] = self.translation()
+        return hmat
+
+    @classmethod
+    def exp(cls, tangent: np.ndarray) -> SE3:
+        assert tangent.shape == (SE3.tangent_dim,)
+        rotation = SO3.exp(tangent[3:])
+        theta_squared = tangent[3:] @ tangent[3:]
+        use_taylor = theta_squared < get_epsilon(theta_squared.dtype)
+        theta_squared_safe = 1.0 if use_taylor else theta_squared
+        theta_safe = np.sqrt(theta_squared_safe)
+        skew_omega = skew(tangent[3:])
+        if use_taylor:
+            V = rotation.as_matrix()
+        else:
+            V = (
+                np.eye(3, dtype=np.float64)
+                + (1.0 - np.cos(theta_safe)) / (theta_squared_safe) * skew_omega
+                + (theta_safe - np.sin(theta_safe))
+                / (theta_squared_safe * theta_safe)
+                * (skew_omega @ skew_omega)
+            )
+        return SE3.from_rotation_and_translation(
+            rotation=rotation,
+            translation=V @ tangent[:3],
+        )
+
+    def inverse(self) -> SE3:
+        R_inv = self.rotation().inverse()
+        return SE3.from_rotation_and_translation(
+            rotation=R_inv,
+            translation=-(R_inv @ self.translation()),
+        )
+
+    def normalize(self) -> SE3:
+        return SE3.from_rotation_and_translation(
+            rotation=self.rotation().normalize(),
+            translation=self.translation(),
+        )
+
+    def apply(self, target: np.ndarray) -> np.ndarray:
+        assert target.shape == (SE3.space_dim,)
+        return self.rotation() @ target + self.translation()
+
+    def multiply(self, other: SE3) -> SE3:
+        return SE3.from_rotation_and_translation(
+            rotation=self.rotation() @ other.rotation(),
+            translation=(self.rotation() @ other.translation()) + self.translation(),
+        )
+
+    def log(self) -> np.ndarray:
+        omega = self.rotation().log()
+        theta_squared = omega @ omega
+        use_taylor = theta_squared < get_epsilon(theta_squared.dtype)
+        skew_omega = skew(omega)
+        theta_squared_safe = 1.0 if use_taylor else theta_squared
+        theta_safe = np.sqrt(theta_squared_safe)
+        half_theta_safe = 0.5 * theta_safe
+        skew_omega_norm = skew_omega @ skew_omega
+        if use_taylor:
+            V_inv = (
+                np.eye(3, dtype=np.float64) - 0.5 * skew_omega + skew_omega_norm / 12.0
+            )
+        else:
+            V_inv = (
+                np.eye(3, dtype=np.float64)
+                - 0.5 * skew_omega
+                + (
+                    1.0
+                    - theta_safe
+                    * np.cos(half_theta_safe)
+                    / (2.0 * np.sin(half_theta_safe))
+                )
+                / theta_squared_safe
+                * skew_omega_norm
+            )
+        return np.concatenate([V_inv @ self.translation(), omega])
+
+    def adjoint(self) -> np.ndarray:
+        R = self.rotation().as_matrix()
+        return np.block(
+            [
+                [R, skew(self.translation()) @ R],
+                [np.zeros((3, 3), dtype=np.float64), R],
+            ]
+        )
+
+    # Jacobians.
+
+    # Eqn 179 a)
+    @classmethod
+    def ljac(cls, other: np.ndarray) -> np.ndarray:
+        theta = other[3:]
+        if theta @ theta < get_epsilon(theta.dtype):
+            return np.eye(cls.tangent_dim)
+        Q = _getQ(other)
+        J = SO3.ljac(theta)
+        O = np.zeros((3, 3))
+        return np.block([[J, Q], [O, J]])
+
+    # Eqn 179 b)
+    @classmethod
+    def ljacinv(cls, other: np.ndarray) -> np.ndarray:
+        theta = other[3:]
+        if theta @ theta < get_epsilon(theta.dtype):
+            return np.eye(cls.tangent_dim)
+        Q = _getQ(other)
+        J_inv = SO3.ljacinv(theta)
+        O = np.zeros((3, 3))
+        return np.block([[J_inv, -J_inv @ Q @ J_inv], [O, J_inv]])
+
+
+# Eqn 180.
+def _getQ(c) -> np.ndarray:
+    theta_sq = c[3:] @ c[3:]
+    A = 0.5
+    if theta_sq < get_epsilon(theta_sq.dtype):
+        B = (1.0 / 6.0) + (1.0 / 120.0) * theta_sq
+        C = -(1.0 / 24.0) + (1.0 / 720.0) * theta_sq
+        D = -(1.0 / 60.0)
+    else:
+        theta = np.sqrt(theta_sq)
+        sin_theta = np.sin(theta)
+        cos_theta = np.cos(theta)
+        B = (theta - sin_theta) / (theta_sq * theta)
+        C = (1.0 - theta_sq / 2.0 - cos_theta) / (theta_sq * theta_sq)
+        D = ((2) * theta - (3) * sin_theta + theta * cos_theta) / (
+            (2) * theta_sq * theta_sq * theta
+        )
+    V = skew(c[:3])
+    W = skew(c[3:])
+    VW = V @ W
+    WV = VW.T
+    WVW = WV @ W
+    VWW = VW @ W
+    return (
+        +A * V
+        + B * (WV + VW + WVW)
+        - C * (VWW - VWW.T - 3 * WVW)
+        + D * (WVW @ W + W @ WVW)
+    )

mink/lie/so3.py

@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import mujoco
+import numpy as np
+
+from .base import MatrixLieGroup
+from .utils import get_epsilon, skew
+
+_IDENTITIY_WXYZ = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
+_INVERT_QUAT_SIGN = np.array([1.0, -1.0, -1.0, -1.0], dtype=np.float64)
+
+
+class RollPitchYaw(NamedTuple):
+    roll: float
+    pitch: float
+    yaw: float
+
+
+@dataclass(frozen=True)
+class SO3(MatrixLieGroup):
+    """Special orthogonal group for 3D rotations.
+
+    Internal parameterization is (qw, qx, qy, qz). Tangent parameterization is
+    (omega_x, omega_y, omega_z).
+    """
+
+    wxyz: np.ndarray
+    matrix_dim: int = 3
+    parameters_dim: int = 4
+    tangent_dim: int = 3
+    space_dim: int = 3
+
+    def __post_init__(self) -> None:
+        if self.wxyz.shape != (self.parameters_dim,):
+            raise ValueError(
+                f"Expeced wxyz to be a length 4 vector but got {self.wxyz.shape[0]}."
+            )
+
+    def __repr__(self) -> str:
+        wxyz = np.round(self.wxyz, 5)
+        return f"{self.__class__.__name__}(wxyz={wxyz})"
+
+    def parameters(self) -> np.ndarray:
+        return self.wxyz
+
+    def copy(self) -> SO3:
+        return SO3(wxyz=self.wxyz.copy())
+
+    @classmethod
+    def from_x_radians(cls, theta: float) -> SO3:
+        return SO3.exp(np.array([theta, 0.0, 0.0], dtype=np.float64))
+
+    @classmethod
+    def from_y_radians(cls, theta: float) -> SO3:
+        return SO3.exp(np.array([0.0, theta, 0.0], dtype=np.float64))
+
+    @classmethod
+    def from_z_radians(cls, theta: float) -> SO3:
+        return SO3.exp(np.array([0.0, 0.0, theta], dtype=np.float64))
+
+    @classmethod
+    def from_rpy_radians(
+        cls,
+        roll: float,
+        pitch: float,
+        yaw: float,
+    ) -> SO3:
+        return (
+            SO3.from_z_radians(yaw)
+            @ SO3.from_y_radians(pitch)
+            @ SO3.from_x_radians(roll)
+        )
+
+    @classmethod
+    def from_matrix(cls, matrix: np.ndarray) -> SO3:
+        assert matrix.shape == (SO3.matrix_dim, SO3.matrix_dim)
+        wxyz = np.zeros(SO3.parameters_dim, dtype=np.float64)
+        mujoco.mju_mat2Quat(wxyz, matrix.ravel())
+        return SO3(wxyz=wxyz)
+
+    @classmethod
+    def identity(cls) -> SO3:
+        return SO3(wxyz=_IDENTITIY_WXYZ)
+
+    @classmethod
+    def sample_uniform(cls) -> SO3:
+        # Ref: https://lavalle.pl/planning/node198.html
+        u1, u2, u3 = np.random.uniform(
+            low=np.zeros(shape=(3,)),
+            high=np.array([1.0, 2.0 * np.pi, 2.0 * np.pi]),
+        )
+        a = np.sqrt(1.0 - u1)
+        b = np.sqrt(u1)
+        wxyz = np.array(
+            [
+                a * np.sin(u2),
+                a * np.cos(u2),
+                b * np.sin(u3),
+                b * np.cos(u3),
+            ],
+            dtype=np.float64,
+        )
+        return SO3(wxyz=wxyz)
+
+    # Eq. 138.
+    def as_matrix(self) -> np.ndarray:
+        mat = np.zeros(9, dtype=np.float64)
+        mujoco.mju_quat2Mat(mat, self.wxyz)
+        return mat.reshape(3, 3)
+
+    def compute_roll_radians(self) -> float:
+        q0, q1, q2, q3 = self.wxyz
+        return np.arctan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1**2 + q2**2))
+
+    def compute_pitch_radians(self) -> float:
+        q0, q1, q2, q3 = self.wxyz
+        return np.arcsin(2 * (q0 * q2 - q3 * q1))
+
+    def compute_yaw_radians(self) -> float:
+        q0, q1, q2, q3 = self.wxyz
+        return np.arctan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2**2 + q3**2))
+
+    def as_rpy_radians(self) -> RollPitchYaw:
+        return RollPitchYaw(
+            roll=self.compute_roll_radians(),
+            pitch=self.compute_pitch_radians(),
+            yaw=self.compute_yaw_radians(),
+        )
+
+    # Paragraph above Appendix B.A.
+    def inverse(self) -> SO3:
+        return SO3(wxyz=self.wxyz * _INVERT_QUAT_SIGN)
+
+    def normalize(self) -> SO3:
+        return SO3(wxyz=self.wxyz / np.linalg.norm(self.wxyz))
+
+    # Eq. 136.
+    def apply(self, target: np.ndarray) -> np.ndarray:
+        assert target.shape == (SO3.space_dim,)
+        padded_target = np.concatenate([np.zeros(1, dtype=np.float64), target])
+        return (self @ SO3(wxyz=padded_target) @ self.inverse()).wxyz[1:]
+
+    def multiply(self, other: SO3) -> SO3:
+        res = np.empty(self.parameters_dim, dtype=np.float64)
+        mujoco.mju_mulQuat(res, self.wxyz, other.wxyz)
+        return SO3(wxyz=res)
+
+    ##
+    #
+    ##
+
+    # Eq. 132.
+    @classmethod
+    def exp(cls, tangent: np.ndarray) -> SO3:
+        assert tangent.shape == (SO3.tangent_dim,)
+        theta_squared = tangent @ tangent
+        theta_pow_4 = theta_squared * theta_squared
+        use_taylor = theta_squared < get_epsilon(tangent.dtype)
+        safe_theta = 1.0 if use_taylor else np.sqrt(theta_squared)
+        safe_half_theta = 0.5 * safe_theta
+        if use_taylor:
+            real = 1.0 - theta_squared / 8.0 + theta_pow_4 / 384.0
+            imaginary = 0.5 - theta_squared / 48.0 + theta_pow_4 / 3840.0
+        else:
+            real = np.cos(safe_half_theta)
+            imaginary = np.sin(safe_half_theta) / safe_theta
+        wxyz = np.concatenate([np.array([real]), imaginary * tangent])
+        return SO3(wxyz=wxyz)
+
+    # Eq. 133.
+    def log(self) -> np.ndarray:
+        w = self.wxyz[0]
+        norm_sq = self.wxyz[1:] @ self.wxyz[1:]
+        use_taylor = norm_sq < get_epsilon(norm_sq.dtype)
+        norm_safe = 1.0 if use_taylor else np.sqrt(norm_sq)
+        w_safe = w if use_taylor else 1.0
+        atan_n_over_w = np.arctan2(-norm_safe if w < 0 else norm_safe, abs(w))
+        if use_taylor:
+            atan_factor = 2.0 / w_safe - 2.0 / 3.0 * norm_sq / w_safe**3
+        else:
+            if abs(w) < get_epsilon(w.dtype):
+                scl = 1.0 if w > 0.0 else -1.0
+                atan_factor = scl * np.pi / norm_safe
+            else:
+                atan_factor = 2.0 * atan_n_over_w / norm_safe
+        return atan_factor * self.wxyz[1:]
+
+    # Eq. 139.
+    def adjoint(self) -> np.ndarray:
+        return self.as_matrix()
+
+    # Jacobians.
+
+    # Eqn. 145, 174.
+    @classmethod
+    def ljac(cls, other: np.ndarray) -> np.ndarray:
+        theta = np.sqrt(other @ other)
+        use_taylor = theta < get_epsilon(theta.dtype)
+        if use_taylor:
+            t2 = theta**2
+            A = (1.0 / 2.0) * (1.0 - t2 / 12.0 * (1.0 - t2 / 30.0 * (1.0 - t2 / 56.0)))
+            B = (1.0 / 6.0) * (1.0 - t2 / 20.0 * (1.0 - t2 / 42.0 * (1.0 - t2 / 72.0)))
+        else:
+            A = (1 - np.cos(theta)) / (theta**2)
+            B = (theta - np.sin(theta)) / (theta**3)
+        skew_other = skew(other)
+        return np.eye(3) + A * skew_other + B * (skew_other @ skew_other)
+
+    @classmethod
+    def ljacinv(cls, other: np.ndarray) -> np.ndarray:
+        theta = np.sqrt(other @ other)
+        use_taylor = theta < get_epsilon(theta.dtype)
+        if use_taylor:
+            t2 = theta**2
+            A = (1.0 / 12.0) * (1.0 + t2 / 60.0 * (1.0 + t2 / 42.0 * (1.0 + t2 / 40.0)))
+        else:
+            A = (1.0 / theta**2) * (
+                1.0 - (theta * np.sin(theta) / (2.0 * (1.0 - np.cos(theta))))
+            )
+        skew_other = skew(other)
+        return np.eye(3) - 0.5 * skew_other + A * (skew_other @ skew_other)

mink/lie/tests/test_axioms.py

@@ -0,0 +1,51 @@
+"""Tests for group axioms."""
+
+from typing import Type
+
+from absl.testing import absltest, parameterized
+
+from mink import lie
+
+from ..base import MatrixLieGroup
+from .utils import assert_transforms_close
+
+
+@parameterized.named_parameters(
+    ("SO3", lie.SO3),
+    ("SE3", lie.SE3),
+)
+class TestAxioms(parameterized.TestCase):
+    def test_closure(self, group: Type[MatrixLieGroup]):
+        transform_a = group.sample_uniform()
+        transform_b = group.sample_uniform()
+        composed = transform_a @ transform_b
+        assert_transforms_close(composed, composed.normalize())
+        composed = transform_b @ transform_a
+        assert_transforms_close(composed, composed.normalize())
+        composed = transform_a @ transform_b
+        assert_transforms_close(composed, composed.normalize())
+
+    def test_identity(self, group: Type[MatrixLieGroup]):
+        transform = group.sample_uniform()
+        identity = group.identity()
+        assert_transforms_close(transform, identity @ transform)
+        assert_transforms_close(transform, transform @ identity)
+
+    def test_inverse(self, group: Type[MatrixLieGroup]):
+        transform = group.sample_uniform()
+        identity = group.identity()
+        assert_transforms_close(identity, transform.inverse() @ transform)
+        assert_transforms_close(identity, transform @ transform.inverse())
+
+    def test_associative(self, group: Type[MatrixLieGroup]):
+        transform_a = group.sample_uniform()
+        transform_b = group.sample_uniform()
+        transform_c = group.sample_uniform()
+        assert_transforms_close(
+            (transform_a @ transform_b) @ transform_c,
+            transform_a @ (transform_b @ transform_c),
+        )
+
+
+if __name__ == "__main__":
+    absltest.main()

mink/lie/tests/test_operations.py

@@ -0,0 +1,77 @@
+"""Tests for general operation definitions."""
+
+from typing import Type
+
+import numpy as np
+from absl.testing import absltest, parameterized
+
+from ..base import MatrixLieGroup
+from ..se3 import SE3
+from ..so3 import SO3
+from .utils import assert_transforms_close
+
+
+@parameterized.named_parameters(
+    ("SO3", SO3),
+    ("SE3", SE3),
+)
+class TestOperations(parameterized.TestCase):
+    def test_inverse_bijective(self, group: Type[MatrixLieGroup]):
+        """Check inverse of inverse."""
+        transform = group.sample_uniform()
+        assert_transforms_close(transform, transform.inverse().inverse())
+
+    def test_matrix_bijective(self, group: Type[MatrixLieGroup]):
+        """Check that we can convert to and from matrices."""
+        transform = group.sample_uniform()
+        assert_transforms_close(transform, group.from_matrix(transform.as_matrix()))
+
+    def test_log_exp_bijective(self, group: Type[MatrixLieGroup]):
+        """Check 1-to-1 mapping for log <=> exp operations."""
+        transform = group.sample_uniform()
+
+        tangent = transform.log()
+        self.assertEqual(tangent.shape, (group.tangent_dim,))
+
+        exp_transform = group.exp(tangent)
+        assert_transforms_close(transform, exp_transform)
+        np.testing.assert_allclose(tangent, exp_transform.log())
+
+    def test_adjoint(self, group: Type[MatrixLieGroup]):
+        transform = group.sample_uniform()
+        omega = np.random.randn(group.tangent_dim)
+        assert_transforms_close(
+            transform @ group.exp(omega),
+            group.exp(transform.adjoint() @ omega) @ transform,
+        )
+
+    def test_rminus(self, group: Type[MatrixLieGroup]):
+        T_wa = group.sample_uniform()
+        T_wb = group.sample_uniform()
+        T_ab = T_wa.inverse() @ T_wb
+        np.testing.assert_allclose(T_wb.rminus(T_wa), T_ab.log())
+
+    def test_rplus(self, group: Type[MatrixLieGroup]):
+        T_wa = group.sample_uniform()
+        T_wb = group.sample_uniform()
+        T_ab = T_wa.inverse() @ T_wb
+        assert_transforms_close(T_wa.rplus(T_ab.log()), T_wb)
+
+    def test_jlog(self, group: Type[MatrixLieGroup]):
+        state = group.sample_uniform()
+        w = np.random.rand(state.tangent_dim) * 1e-4
+        state_pert = state.plus(w).log()
+        state_lin = state.log() + state.jlog() @ w
+        np.testing.assert_allclose(state_pert, state_lin, atol=1e-7)
+
+
+class TestGroupSpecificOperations(absltest.TestCase):
+    """Group specific tests."""
+
+    def test_so3_rpy_bijective(self):
+        T = SO3.sample_uniform()
+        assert_transforms_close(T, SO3.from_rpy_radians(*T.as_rpy_radians()))
+
+
+if __name__ == "__main__":
+    absltest.main()

mink/lie/tests/utils.py

@@ -0,0 +1,20 @@
+import numpy as np
+
+from ..base import MatrixLieGroup
+from ..se3 import SE3, SO3
+
+
+def assert_transforms_close(a: MatrixLieGroup, b: MatrixLieGroup) -> None:
+    np.testing.assert_allclose(a.as_matrix(), b.as_matrix(), atol=1e-7)
+
+    # Account for quaternion double cover (q = -q).
+    pa = a.parameters()
+    pb = b.parameters()
+    if isinstance(a, SO3):
+        pa *= np.sign(pa[0])
+        pb *= np.sign(pb[0])
+    elif isinstance(a, SE3):
+        pa[:4] *= np.sign(pa[0])
+        pb[:4] *= np.sign(pb[0])
+
+    np.testing.assert_allclose(pa, pb, atol=1e-7)

mink/lie/utils.py

@@ -0,0 +1,29 @@
+import mujoco
+import numpy as np
+
+
+def get_epsilon(dtype: np.dtype) -> float:
+    return {
+        np.dtype("float32"): 1e-5,
+        np.dtype("float64"): 1e-10,
+    }[dtype]
+
+
+def skew(x: np.ndarray) -> np.ndarray:
+    assert x.shape == (3,)
+    wx, wy, wz = x
+    return np.array(
+        [
+            [0.0, -wz, wy],
+            [wz, 0.0, -wx],
+            [-wy, wx, 0.0],
+        ]
+    )
+
+
+def mat2quat(mat: np.ndarray):
+    """Convert a MuJoCo matrix (9,) to a quaternion (4,)."""
+    assert mat.shape == (9,)
+    quat = np.empty(4, dtype=np.float64)
+    mujoco.mju_mat2Quat(quat, mat)
+    return quat

mink/limits/__init__.py

@@ -0,0 +1,16 @@
+"""Kinematic limits."""
+
+from .collision_avoidance_limit import CollisionAvoidanceLimit
+from .configuration_limit import ConfigurationLimit
+from .exceptions import LimitDefinitionError
+from .limit import Constraint, Limit
+from .velocity_limit import VelocityLimit
+
+__all__ = (
+    "ConfigurationLimit",
+    "CollisionAvoidanceLimit",
+    "Constraint",
+    "Limit",
+    "VelocityLimit",
+    "LimitDefinitionError",
+)