mink 0.0.8__py3-none-any.whl → 0.0.9__py3-none-any.whl

mink/lie/se3.py

@@ -113,130 +113,131 @@ class SE3(MatrixLieGroup):
 
     @classmethod
     def exp(cls, tangent: np.ndarray) -> SE3:
-        assert tangent.shape == (SE3.tangent_dim,)
+        assert tangent.shape == (cls.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()
+        theta = np.float64(mujoco.mju_norm3(tangent[3:]))
+        t2 = theta * theta
+        if t2 < get_epsilon(t2.dtype):
+            v_mat = rotation.as_matrix()
         else:
-            V = (
+            skew_omega = skew(tangent[3:])
+            v_mat = (
                 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)
+                + (1.0 - np.cos(theta)) / t2 * skew_omega
+                + (theta - np.sin(theta)) / (t2 * theta) * (skew_omega @ skew_omega)
             )
-        return SE3.from_rotation_and_translation(
+        return cls.from_rotation_and_translation(
             rotation=rotation,
-            translation=V @ tangent[:3],
+            translation=v_mat @ 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()),
+        inverse_wxyz_xyz = np.empty(SE3.parameters_dim, dtype=np.float64)
+        mujoco.mju_negQuat(inverse_wxyz_xyz[:4], self.wxyz_xyz[:4])
+        mujoco.mju_rotVecQuat(
+            inverse_wxyz_xyz[4:], -1.0 * self.wxyz_xyz[4:], inverse_wxyz_xyz[:4]
         )
+        return SE3(wxyz_xyz=inverse_wxyz_xyz)
 
     def normalize(self) -> SE3:
-        return SE3.from_rotation_and_translation(
-            rotation=self.rotation().normalize(),
-            translation=self.translation(),
-        )
+        normalized_wxyz_xyz = np.array(self.wxyz_xyz)
+        mujoco.mju_normalize4(normalized_wxyz_xyz[:4])
+        return SE3(wxyz_xyz=normalized_wxyz_xyz)
 
     def apply(self, target: np.ndarray) -> np.ndarray:
         assert target.shape == (SE3.space_dim,)
-        return self.rotation() @ target + self.translation()
+        rotated_target = np.empty(SE3.space_dim, dtype=np.float64)
+        mujoco.mju_rotVecQuat(rotated_target, target, self.wxyz_xyz[:4])
+        return rotated_target + self.wxyz_xyz[4:]
 
     def multiply(self, other: SE3) -> SE3:
-        return SE3.from_rotation_and_translation(
-            rotation=self.rotation() @ other.rotation(),
-            translation=(self.rotation() @ other.translation()) + self.translation(),
-        )
+        wxyz_xyz = np.empty(SE3.parameters_dim, dtype=np.float64)
+        mujoco.mju_mulQuat(wxyz_xyz[:4], self.wxyz_xyz[:4], other.wxyz_xyz[:4])
+        mujoco.mju_rotVecQuat(wxyz_xyz[4:], other.wxyz_xyz[4:], self.wxyz_xyz[:4])
+        wxyz_xyz[4:] += self.wxyz_xyz[4:]
+        return SE3(wxyz_xyz=wxyz_xyz)
 
     def log(self) -> np.ndarray:
         omega = self.rotation().log()
-        theta_squared = omega @ omega
-        use_taylor = theta_squared < get_epsilon(theta_squared.dtype)
+        theta = np.float64(mujoco.mju_norm3(omega))
+        t2 = theta * theta
         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
+        skew_omega2 = skew_omega @ skew_omega
+        if t2 < get_epsilon(t2.dtype):
+            vinv_mat = (
+                np.eye(3, dtype=np.float64) - 0.5 * skew_omega + skew_omega2 / 12.0
             )
         else:
-            V_inv = (
+            half_theta = 0.5 * theta
+            vinv_mat = (
                 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
+                + (1.0 - 0.5 * theta * np.cos(half_theta) / np.sin(half_theta))
+                / t2
+                * skew_omega2
             )
-        return np.concatenate([V_inv @ self.translation(), omega])
+        tangent = np.empty(SE3.tangent_dim, dtype=np.float64)
+        tangent[:3] = vinv_mat @ self.translation()
+        tangent[3:] = omega
+        return tangent
 
     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],
-            ]
-        )
+        rotation = self.rotation()
+        rotation_mat = rotation.as_matrix()
+        tangent_mat = skew(self.translation()) @ rotation_mat
+        adjoint_mat = np.zeros((SE3.tangent_dim, SE3.tangent_dim), dtype=np.float64)
+        adjoint_mat[:3, :3] = rotation_mat
+        adjoint_mat[:3, 3:] = tangent_mat
+        adjoint_mat[3:, 3:] = rotation_mat
+        return adjoint_mat
 
     # Jacobians.
 
     # Eqn 179 a)
     @classmethod
     def ljac(cls, other: np.ndarray) -> np.ndarray:
-        theta = other[3:]
-        if theta @ theta < get_epsilon(theta.dtype):
+        theta_squared = np.float64(mujoco.mju_dot3(other[3:], other[3:]))
+        if theta_squared < get_epsilon(theta_squared.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]])
+        ljac_se3 = np.zeros((cls.tangent_dim, cls.tangent_dim), dtype=np.float64)
+        ljac_translation = _getQ(other)
+        ljac_so3 = SO3.ljac(other[3:])
+        ljac_se3[:3, :3] = ljac_so3
+        ljac_se3[:3, 3:] = ljac_translation
+        ljac_se3[3:, 3:] = ljac_so3
+        return ljac_se3
 
     # Eqn 179 b)
     @classmethod
     def ljacinv(cls, other: np.ndarray) -> np.ndarray:
-        theta = other[3:]
-        if theta @ theta < get_epsilon(theta.dtype):
+        theta_squared = np.float64(mujoco.mju_dot3(other[3:], other[3:]))
+        if theta_squared < get_epsilon(theta_squared.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]])
+        ljacinv_se3 = np.zeros((cls.tangent_dim, cls.tangent_dim), dtype=np.float64)
+        ljac_translation = _getQ(other)
+        ljacinv_so3 = SO3.ljacinv(other[3:])
+        ljacinv_se3[:3, :3] = ljacinv_so3
+        ljacinv_se3[:3, 3:] = -ljacinv_so3 @ ljac_translation @ ljacinv_so3
+        ljacinv_se3[3:, 3:] = ljacinv_so3
+        return ljacinv_se3
 
 
 # Eqn 180.
 def _getQ(c) -> np.ndarray:
-    theta_sq = c[3:] @ c[3:]
+    theta = np.float64(mujoco.mju_norm3(c[3:]))
+    t2 = theta * theta
     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
+    if t2 < get_epsilon(t2.dtype):
+        B = (1.0 / 6.0) + (1.0 / 120.0) * t2
+        C = -(1.0 / 24.0) + (1.0 / 720.0) * t2
         D = -(1.0 / 60.0)
     else:
-        theta = np.sqrt(theta_sq)
+        t4 = t2 * t2
         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
-        )
+        B = (theta - sin_theta) / (t2 * theta)
+        C = (1.0 - 0.5 * t2 - cos_theta) / t4
+        D = (2.0 * theta - 3.0 * sin_theta + theta * cos_theta) / (2.0 * t4 * theta)
     V = skew(c[:3])
     W = skew(c[3:])
     VW = V @ W
@@ -246,6 +247,6 @@ def _getQ(c) -> np.ndarray:
     return (
         +A * V
         + B * (WV + VW + WVW)
-        - C * (VWW - VWW.T - 3 * WVW)
+        - C * (VWW - VWW.T - 3.0 * WVW)
         + D * (WVW @ W + W @ WVW)
     )

mink/lie/so3.py

@@ -7,10 +7,9 @@ import mujoco
 import numpy as np
 
 from .base import MatrixLieGroup
-from .utils import get_epsilon, skew
+from .utils import get_epsilon
 
 _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):
@@ -137,63 +136,47 @@ class SO3(MatrixLieGroup):
             yaw=self.compute_yaw_radians(),
         )
 
-    # Paragraph above Appendix B.A.
     def inverse(self) -> SO3:
-        return SO3(wxyz=self.wxyz * _INVERT_QUAT_SIGN)
+        conjugate_wxyz = np.empty(4)
+        mujoco.mju_negQuat(conjugate_wxyz, self.wxyz)
+        return SO3(wxyz=conjugate_wxyz)
 
     def normalize(self) -> SO3:
-        return SO3(wxyz=self.wxyz / np.linalg.norm(self.wxyz))
+        normalized_wxyz = np.array(self.wxyz)
+        mujoco.mju_normalize4(normalized_wxyz)
+        return SO3(wxyz=normalized_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:]
+        rotated_target = np.empty(SO3.space_dim, dtype=np.float64)
+        mujoco.mju_rotVecQuat(rotated_target, target, self.wxyz)
+        return rotated_target
 
     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])
+        axis = np.array(tangent)
+        theta = mujoco.mju_normalize3(axis)
+        wxyz = np.zeros(4)
+        mujoco.mju_axisAngle2Quat(wxyz, axis, theta)
         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
+        if self.wxyz[0] < 0.0:
+            theta = 2.0 * np.arccos(-self.wxyz[0])
+            axis = -1.0 * np.array(self.wxyz[1:])
         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:]
+            theta = 2.0 * np.arccos(self.wxyz[0])
+            axis = np.array(self.wxyz[1:])
+        mujoco.mju_normalize3(axis)
+        return theta * axis
 
     # Eq. 139.
     def adjoint(self) -> np.ndarray:
@@ -204,28 +187,73 @@ class SO3(MatrixLieGroup):
     # 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)))
+        theta = np.float64(mujoco.mju_norm3(other))
+        t2 = theta * theta
+        if theta < get_epsilon(theta.dtype):
+            alpha = (1.0 / 2.0) * (
+                1.0 - t2 / 12.0 * (1.0 - t2 / 30.0 * (1.0 - t2 / 56.0))
+            )
+            beta = (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)
+            t3 = t2 * theta
+            alpha = (1 - np.cos(theta)) / t2
+            beta = (theta - np.sin(theta)) / t3
+        # ljac = eye(3) + alpha * skew_other + beta * (skew_other @ skew_other)
+        ljac = np.empty((3, 3))
+        # skew_other @ skew_other == outer(other) - inner(other) * eye(3)
+        mujoco.mju_mulMatMat(ljac, other.reshape(3, 1), other.reshape(1, 3))
+        inner_product = mujoco.mju_dot3(other, other)
+        ljac[0, 0] -= inner_product
+        ljac[1, 1] -= inner_product
+        ljac[2, 2] -= inner_product
+        ljac *= beta
+        # + alpha * skew_other
+        alpha_vec = alpha * other
+        ljac[0, 1] += -alpha_vec[2]
+        ljac[0, 2] += alpha_vec[1]
+        ljac[1, 0] += alpha_vec[2]
+        ljac[1, 2] += -alpha_vec[0]
+        ljac[2, 0] += -alpha_vec[1]
+        ljac[2, 1] += alpha_vec[0]
+        # + eye(3)
+        ljac[0, 0] += 1.0
+        ljac[1, 1] += 1.0
+        ljac[2, 2] += 1.0
+        return ljac
 
     @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)))
+        theta = np.float64(mujoco.mju_norm3(other))
+        t2 = theta * theta
+        if theta < get_epsilon(theta.dtype):
+            beta = (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) * (
+            beta = (1.0 / t2) * (
                 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)
+        # ljacinv = eye(3) - 0.5 * skew_other + beta * (skew_other @ skew_other)
+        ljacinv = np.empty((3, 3))
+        # skew_other @ skew_other == outer(other) - inner(other) * eye(3)
+        mujoco.mju_mulMatMat(ljacinv, other.reshape(3, 1), other.reshape(1, 3))
+        inner_product = mujoco.mju_dot3(other, other)
+        ljacinv[0, 0] -= inner_product
+        ljacinv[1, 1] -= inner_product
+        ljacinv[2, 2] -= inner_product
+        ljacinv *= beta
+        # - 0.5 * skew_other
+        alpha_vec = -0.5 * other
+        ljacinv[0, 1] += -alpha_vec[2]
+        ljacinv[0, 2] += alpha_vec[1]
+        ljacinv[1, 0] += alpha_vec[2]
+        ljacinv[1, 2] += -alpha_vec[0]
+        ljacinv[2, 0] += -alpha_vec[1]
+        ljacinv[2, 1] += alpha_vec[0]
+        # + eye(3)
+        ljacinv[0, 0] += 1.0
+        ljacinv[1, 1] += 1.0
+        ljacinv[2, 2] += 1.0
+        return ljacinv

{mink-0.0.8.dist-info → mink-0.0.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mink
-Version: 0.0.8
+Version: 0.0.9
 Summary: mink: MuJoCo inverse kinematics.
 Keywords: inverse,kinematics,mujoco
 Author-email: Kevin Zakka <zakka@berkeley.edu>
@@ -21,7 +21,7 @@ License-File: LICENSE
 Requires-Dist: mujoco >= 3.1.6
 Requires-Dist: qpsolvers[daqp] >= 4.3.1
 Requires-Dist: typing_extensions
-Requires-Dist: numpy < 2.0.0
+Requires-Dist: numpy
 Requires-Dist: mink[examples, dev] ; extra == "all"
 Requires-Dist: black ; extra == "dev"
 Requires-Dist: mink[test] ; extra == "dev"

{mink-0.0.8.dist-info → mink-0.0.9.dist-info}/RECORD

@@ -333,8 +333,8 @@ mink/contrib/keyboard_teleop/keycodes.py,sha256=1AN3bWpkw99ZjrkIQnu8QdHSz-_boWMY
 mink/contrib/keyboard_teleop/teleop_mocap.py,sha256=2cuauiBHsQb6IVruRbGtCumdHQazEGvt9w2L50eKGzw,8095
 mink/lie/__init__.py,sha256=7tm3ZFnF3o1SDd9MOFO1In13lHMQJHO0FB-ejI6tsgE,202
 mink/lie/base.py,sha256=8TGWWEWRvmoBsJP-9y1YXQHkTC-TFfI8JDvBFKxxP-A,4865
-mink/lie/se3.py,sha256=mmPjbljB0ddPrV0b47OV8GP9rqU6f-bpW8TQPW_p-9A,8050
-mink/lie/so3.py,sha256=IqXqalM5NG0cTi1ZZcuXA2lcAk7A3Jnt_1O2n-ae0DE,7474
+mink/lie/se3.py,sha256=cEzAzKEfBAQcbHhbZbhW8QZA8WpfRRaeZ0YGdSmhOx8,8634
+mink/lie/so3.py,sha256=0KgvflDlUJHd_Ruy2pDlgyR7ZC1HkJz4HVxWcCLo5EY,8182
 mink/lie/utils.py,sha256=DuEl3pj84daLvMKN84-YBvkXnJfqvc5vpvJ9J-pJ11U,403
 mink/limits/__init__.py,sha256=hX5Dgpri9AE6YtUCkW79AXMBuNAuBhniR9kQ6Rxwv3Y,416
 mink/limits/collision_avoidance_limit.py,sha256=MJXesSjvm_2O6RUEz2e4D8uTVOHyTsH9j4UxkPemcn8,11918
@@ -351,7 +351,7 @@ mink/tasks/frame_task.py,sha256=DeNg-bEJxqnrRI0FaJK4UHyb8CyF4A7uPF9G-CdN0sg,5307
 mink/tasks/posture_task.py,sha256=sVDZyalCh5DP8zmCuX5kYuZxiMxRut_mTZUjv5y3b5M,3998
 mink/tasks/relative_frame_task.py,sha256=9rIiYocI1eEESt0bLZZpQR7K6ggVRZH8iEhdhzkfZa0,5119
 mink/tasks/task.py,sha256=F16YZT1L9ueNOcKoOhbCyEnZw0DOgrmjqADl0jahVQI,4838
-mink-0.0.8.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mink-0.0.8.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-mink-0.0.8.dist-info/METADATA,sha256=SoXK-YNeZsu55hyaqB2hPAPjtWXR7xH7wL9cc2vBTBU,5614
-mink-0.0.8.dist-info/RECORD,,
+mink-0.0.9.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mink-0.0.9.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+mink-0.0.9.dist-info/METADATA,sha256=PbX9fwluNc1-M5fB8cuJnPD33Flma45GzAU5yd6myfg,5606
+mink-0.0.9.dist-info/RECORD,,