reachy-mini 1.0.0__py3-none-any.whl

reachy_mini/kinematics/analytical_kinematics.py

@@ -0,0 +1,102 @@
+"""An analytical kinematics engine for Reachy Mini, using Rust bindings.
+
+The inverse kinematics use an analytical method, while the forward kinematics
+use a numerical method (Newton).
+"""
+
+import json
+from importlib.resources import files
+from typing import Annotated
+
+import numpy as np
+from numpy.typing import NDArray
+from reachy_mini_rust_kinematics import ReachyMiniRustKinematics
+
+import reachy_mini
+
+# Duplicated for now.
+SLEEP_HEAD_POSE = np.array(
+    [
+        [0.911, 0.004, 0.413, -0.021],
+        [-0.004, 1.0, -0.001, 0.001],
+        [-0.413, -0.001, 0.911, -0.044],
+        [0.0, 0.0, 0.0, 1.0],
+    ]
+)
+
+
+class AnalyticalKinematics:
+    """Reachy Mini Analytical Kinematics class, implemented in Rust with python bindings."""
+
+    def __init__(self) -> None:
+        """Initialize."""
+        assets_root_path: str = str(files(reachy_mini).joinpath("assets/"))
+        data_path = assets_root_path + "/kinematics_data.json"
+        data = json.load(open(data_path, "rb"))
+
+        self.head_z_offset = data["head_z_offset"]
+
+        self.kin = ReachyMiniRustKinematics(
+            data["motor_arm_length"], data["rod_length"]
+        )
+
+        self.start_body_yaw = 0.0
+
+        self.motors = data["motors"]
+        for motor in self.motors:
+            self.kin.add_branch(
+                motor["branch_position"],
+                np.linalg.inv(motor["T_motor_world"]),  # type: ignore[arg-type]
+                1 if motor["solution"] else -1,
+            )
+
+        sleep_head_pose = SLEEP_HEAD_POSE.copy()
+        sleep_head_pose[:3, 3][2] += self.head_z_offset
+        self.kin.reset_forward_kinematics(sleep_head_pose)  # type: ignore[arg-type]
+
+    def ik(
+        self,
+        pose: Annotated[NDArray[np.float64], (4, 4)],
+        body_yaw: float = 0.0,
+        check_collision: bool = False,
+        no_iterations: int = 0,
+    ) -> Annotated[NDArray[np.float64], (7,)]:
+        """Compute the inverse kinematics for a given head pose.
+
+        check_collision and no_iterations are not used by AnalyticalKinematics. We keep them for compatibility with the other kinematics engines
+        """
+        _pose = pose.copy()
+        _pose[:3, 3][2] += self.head_z_offset
+
+        stewart_joints = self.kin.inverse_kinematics(_pose, body_yaw)  # type: ignore[arg-type]
+
+        return np.array([body_yaw] + stewart_joints)
+
+    def fk(
+        self,
+        joint_angles: Annotated[NDArray[np.float64], (7,)],
+        check_collision: bool = False,
+        no_iterations: int = 3,
+    ) -> Annotated[NDArray[np.float64], (4, 4)]:
+        """Compute the forward kinematics for a given set of joint angles.
+
+        check_collision is not used by AnalyticalKinematics.
+        """
+        body_yaw = joint_angles[0]
+
+        _joint_angles = joint_angles[1:].tolist()
+
+        if no_iterations < 1:
+            raise ValueError("no_iterations must be at least 1")
+
+        T_world_platform = None
+        for _ in range(no_iterations):
+            T_world_platform = np.array(
+                self.kin.forward_kinematics(_joint_angles, body_yaw)
+            )
+
+        assert T_world_platform is not None
+
+        T_world_platform[:3, 3][2] -= self.head_z_offset
+
+        return T_world_platform

reachy_mini/kinematics/nn_kinematics.py

@@ -0,0 +1,100 @@
+"""Neural Network based FK/IK."""
+
+import time
+from typing import Annotated
+
+import numpy as np
+import numpy.typing as npt
+import onnxruntime
+from scipy.spatial.transform import Rotation as R
+
+
+class NNKinematics:
+    """Neural Network based FK/IK. Fitted from PlacoKinematics data."""
+
+    def __init__(self, models_root_path: str):
+        """Intialize."""
+        self.fk_model_path = f"{models_root_path}/fknetwork.onnx"
+        self.ik_model_path = f"{models_root_path}/iknetwork.onnx"
+        self.fk_infer = OnnxInfer(self.fk_model_path)
+        self.ik_infer = OnnxInfer(self.ik_model_path)
+
+        self.start_body_yaw = 0.0  # No used, kept for compatibility
+
+    def ik(
+        self,
+        pose: Annotated[npt.NDArray[np.float64], (4, 4)],
+        body_yaw: float = 0.0,
+        check_collision: bool = False,
+        no_iterations: int = 0,
+    ) -> Annotated[npt.NDArray[np.float64], (7,)]:
+        """check_collision and no_iterations are not used by NNKinematics.
+
+        We keep them for compatibility with the other kinematics engines
+        """
+        x, y, z = pose[:3, 3][0], pose[:3, 3][1], pose[:3, 3][2]
+        roll, pitch, yaw = R.from_matrix(pose[:3, :3]).as_euler("xyz")
+
+        yaw += body_yaw
+        input = np.array([x, y, z, roll, pitch, yaw])
+
+        joints = self.ik_infer.infer(input)
+        joints[0] += body_yaw
+
+        return joints
+
+    def fk(
+        self,
+        joint_angles: Annotated[npt.NDArray[np.float64], (7,)],
+        check_collision: bool = False,
+        no_iterations: int = 0,
+    ) -> Annotated[npt.NDArray[np.float64], (4, 4)]:
+        """check_collision and no_iterations are not used by NNKinematics.
+
+        We keep them for compatibility with the other kinematics engines
+        """
+        x, y, z, roll, pitch, yaw = self.fk_infer.infer(joint_angles)
+        pose = np.eye(4)
+        pose[:3, 3] = [x, y, z]
+        pose[:3, :3] = R.from_euler("xyz", [roll, pitch, yaw]).as_matrix()
+        return pose
+
+
+class OnnxInfer:
+    """Infer an onnx model."""
+
+    def __init__(self, onnx_model_path: str) -> None:
+        """Initialize."""
+        self.onnx_model_path = onnx_model_path
+        self.ort_session = onnxruntime.InferenceSession(
+            self.onnx_model_path, providers=["CPUExecutionProvider"]
+        )
+
+    def infer(self, input: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+        """Run inference on the input."""
+        outputs = self.ort_session.run(None, {"input": [input]})
+        res: npt.NDArray[np.float64] = outputs[0][0]
+        return res
+
+
+if __name__ == "__main__":
+    nn_kin = NNKinematics(
+        "assets/models",
+    )
+
+    times_fk: list[float] = []
+    times_ik: list[float] = []
+    for i in range(1000):
+        fk_input = np.random.random(7).astype(np.float64)
+        # ik_input = np.random.random(6).astype(np.float64)
+
+        fk_s = time.time()
+        fk_output = nn_kin.fk(fk_input)
+        times_fk.append(time.time() - fk_s)
+
+        # ik_s = time.time()
+        # ik_output = nn_kin.ik(ik_input)
+        # times_ik.append(time.time() - ik_s)
+
+    print(f"Average FK inference time: {np.mean(times_fk) * 1e6} µs")
+    # print(f"Average IK inference time: {np.mean(times_ik) * 1e6} µs")