Add kinimatics class, update dm_driver, cleanup README

Author: i2rt-software

README.md

@@ -43,11 +43,11 @@ sh scripts/reset_all_can.sh
 
 ## Test YAM Zero Gravity mode
 
+This enables you to launch the robot in zero gravity mode.
 ```bash
 python i2rt/robots/motor_chain_robot.py --model yam --channel can0 --operation_mode gravity_comp
 ```
 
-
 ## YAM Robot Arm Usage
 
 ### Getting started

doc/set_persist_id_socket_can.md

@@ -1,22 +1,24 @@
-# Socket can is easy to use, but the default devie name is can{idx} and will vary depend on the order the device got connected to the computer. Below is the precedure of setting up persist ID for those socket CAN devices. 
- 
-## For canable devices, goto https://canable.io/updater/ to flush its firmware to candlelight to use socketcan, YAM comes with pre-flushed candlelight firmware.
+# Setting Persistent IDs for SocketCAN Devices
 
+SocketCAN is easy to use, but the default device name is `can{idx}`, which can vary depending on the order in which the device is connected to the computer. Below is the procedure for setting up persistent IDs for these SocketCAN devices.
 
-## Step1: find sysfd paths for can devices
+## For Canable Devices
+Visit [Canable Updater](https://canable.io/updater/) to flash the firmware to candlelight to use with SocketCAN. YAM comes with pre-flashed candlelight firmware.
+
+## Step 1: Find sysfd Paths for CAN Devices
 
 ```shell
 $ ls -l /sys/class/net/can*
 ```
 
-this should  give you something liks
+This should give you an output similar to:
 ```shell
 lrwxrwxrwx 1 root root 0 Jul 15 14:35 /sys/class/net/can0 -> ../../devices/platform/soc/your_can_device/can0
 lrwxrwxrwx 1 root root 0 Jul 15 14:35 /sys/class/net/can1 -> ../../devices/platform/soc/your_can_device/can1
 lrwxrwxrwx 1 root root 0 Jul 15 14:35 /sys/class/net/can2 -> ../../devices/platform/soc/your_can_device/can2
 ```
 
-## Step 2: Use udevadm to Gather Attributes
+## Step 2: Use `udevadm` to Gather Attributes
 ```shell
 udevadm info -a -p /sys/class/net/can0 | grep -i serial
 ```
@@ -32,21 +34,21 @@ SUBSYSTEM=="net", ACTION=="add", ATTRS{serial}=="004E00275548501220373234", NAME
 SUBSYSTEM=="net", ACTION=="add", ATTRS{serial}=="0031005F5548501220373234", NAME="can_follow_r"
 ```
 
-IMPORTANT!!!: Name should start with can (for USB-CAN adapter) or en/eth (for EtherCAT-CAN adapter). And the maximum length limit for a CAN interface name is 13 characters.
+**Important:** The name should start with `can` (for USB-CAN adapters) or `en`/`eth` (for EtherCAT-CAN adapters). The maximum length for a CAN interface name is 13 characters.
 
 ## Step 4: Reload udev Rules
 ```shell
 sudo udevadm control --reload-rules && sudo systemctl restart systemd-udevd && sudo udevadm trigger
 ```
 
-If needed, plug unplug the candevice to make sure the change is effective. 
+Unplug and replug the CAN device to ensure the changes take effect.
 
-run the following command to set up the can device, and you need to run this command after every reboot.
+Run the following command to set up the CAN device, and you need to run this command after every reboot.
 ```
 sudo ip link set up can_right type can bitrate 1000000
 ```
 
-## Step 5: Verify the can device
+## Step 5: Verify the CAN device
 ```shell
 $ ip link show
 1: lo: <LOOPBACK,UP,LOWER_UP> mtu 65536 qdisc noqueue state UNKNOWN mode DEFAULT group default qlen 1000
@@ -61,5 +63,4 @@ $ ip link show
     link/can 
 ```
 
-You can see that this can device got it's name can_right/can_left
-
+You should see that the CAN device is named `can_right`/`can_left`.

i2rt/motor_drivers/dm_driver.py

@@ -9,6 +9,10 @@
 import can
 import numpy as np
 
+# set control frequence
+CONTROL_FREQ = 250
+CONTROL_PERIOD = 1.0 / CONTROL_FREQ  # 4 ms
+
 
 @dataclass
 class MotorConstants:
@@ -33,6 +37,7 @@ class MotorType:
     DM4310V = "DM4310V"
     DM4340 = "DM4340"
     DMH6215 = "DMH6215"
+    DM3507 = "DM3507"
 
     @classmethod
     def get_motor_constants(cls, motor_type: str) -> MotorConstants:
@@ -85,6 +90,15 @@ def get_motor_constants(cls, motor_type: str) -> MotorConstants:
                 TORQUE_MAX=10,
                 TORQUE_MIN=-10,
             )
+        elif motor_type == cls.DM3507:
+            return MotorConstants(
+                POSITION_MAX=12.5,
+                POSITION_MIN=-12.5,
+                VELOCITY_MAX=50,
+                VELOCITY_MIN=-50,
+                TORQUE_MAX=5,
+                TORQUE_MIN=-5,
+            )
         else:
             raise ValueError(f"Motor type '{motor_type}' not recognized.")
 
@@ -484,6 +498,7 @@ def parse_recv_message(self, message: can.Message, motor_type: str) -> FeedbackF
         """
         data = message.data
         error_int = (data[0] & 0xF0) >> 4  # TODO: error code seems incorrect, double check
+
         # convert error into hex
         error_hex = hex(error_int)
         error_message = MotorErrorCode.get_error_message(error_int)
@@ -652,46 +667,32 @@ def start_thread(self) -> None:
             print("waiting for the first state")
 
     def _set_torques_and_update_state(self) -> None:
-        last_time = time.time()
-        iteration_count = 0
-        self.motor_interface.try_receive_message()
+        last_step_time = time.time()
         while self.running:
             try:
-                current_time = time.time()
-                elapsed_time = current_time - last_time
+                # Maintain desired control frequency.
+                while time.time() - last_step_time < CONTROL_PERIOD - 0.001:
+                    time.sleep(0.001)
+                curr_time = time.time()
+                step_time = curr_time - last_step_time
+                last_step_time = curr_time
+                if step_time > 0.005:  # 5 ms
+                    print(f"Warning: Step time {1000 * step_time:.3f} ms in {self.__class__.__name__} control_loop")
+
+                # Update state.
                 with self.command_lock:
                     motor_feedback = self._set_commands(self.commands)
                 with self.state_lock:
                     self.state = motor_feedback
-
-                self._update_absolute_positions(motor_feedback)
-                time.sleep(0.001)
-                iteration_count += 1
-                # Optionally reset counter and timer every second or another fixed interval
-                # to keep the displayed frequency relevant to a recent time window
-                if elapsed_time >= 10.0:
-                    control_frequency = iteration_count / elapsed_time
-
-                    # # Overwrite the current line with the new frequency information
-                    # sys.stdout.write(f"\rControl Frequency: {control_frequency:.2f} Hz")
-                    # sys.stdout.flush()  # Ensure it's displayed
-
-                    # print(f"Control Frequency: {control_frequency:.2f} Hz")
-
-                    # Reset the counter and timer
-                    last_time = current_time
-                    iteration_count = 0
-
+                    self._update_absolute_positions(motor_feedback)
             except Exception as e:
                 print(f"DM Error in control loop: {e}")
                 raise e
 
     def _set_commands(self, commands: List[MotorCmd]) -> List[MotorInfo]:
-        # for joint 0 and joint1
         motor_feedback = []
         for idx, motor_info in enumerate(self.motor_list):
             motor_id, motor_type = motor_info
-            # torque = torques[idx] * self.motor_direction[idx]
             torque = commands[idx].torque * self.motor_direction[idx]
             pos = self._joint_position_sim_to_real_idx(commands[idx].pos, idx)
 
@@ -713,8 +714,6 @@ def _set_commands(self, commands: List[MotorCmd]) -> List[MotorInfo]:
                 raise e
 
             motor_feedback.append(fd_back)
-            # motor_pos.append(fd_back.position)
-            # motor_vel.append(fd_back.velocity * self.motor_direction[idx])
         return motor_feedback
 
     def read_states(self, torques: Optional[np.ndarray] = None) -> List[MotorInfo]:

i2rt/robots/kinematics.py

@@ -0,0 +1,126 @@
+import time
+from typing import List, Optional, Tuple
+
+import mink
+import mujoco
+import numpy as np
+
+
+class Kinematics:
+    def __init__(self, xml_path: str, site_name: Optional[str]):
+        """Initialize the Kinematics object.
+
+        Args:
+            xml_path (str): Path to the MuJoCo XML model file.
+            site_name (Optional[str]): Name of the site for which to compute the forward kinematics.
+        """
+        model = mujoco.MjModel.from_xml_path(xml_path)
+        self._configuration = mink.Configuration(model)
+        self._site_name = site_name
+
+    def fk(self, q: np.ndarray, site_name: Optional[str] = None) -> np.ndarray:
+        """Compute the forward kinematics for the given joint configuration.
+
+        Args:
+            q (np.ndarray): The joint configuration.
+            site_name (Optional[str]): Name of the site for which to compute the forward kinematics.
+                       If not provided, the default site name is used.
+
+        Returns:
+            (np.ndarray): Site frame in world frame. Shape: (4, 4)
+        """
+        self._configuration.update(q)
+        site_name = site_name or self._site_name
+        assert site_name is not None, "site_name must be provided"
+        return self._configuration.get_transform_frame_to_world(site_name, "site").as_matrix()
+
+    def ik(
+        self,
+        target_pose: np.ndarray,
+        site_name: str,
+        init_q: Optional[np.ndarray] = None,
+        limits: Optional[List[mink.Limit]] = None,
+        dt: float = 0.01,
+        solver: str = "quadprog",
+        pos_threshold: float = 1e-4,
+        ori_threshold: float = 1e-4,
+        damping: float = 1e-4,
+        max_iters: int = 200,
+        verbose: bool = False,
+    ) -> Tuple[bool, np.ndarray]:
+        """Differential ik solver, leverging mink.
+
+        Args:
+            target_pose (np.ndarray): The target pose to reach.
+            site_name (str): Name of the desired site.
+            limits (List[mink.Limit]): List of limits to enforce.
+            init_q (Optional[np.ndarray]): Initial joint configuration.
+            dt (float): Integration timestep in [s].
+            solver (str): Quadratic program solver.
+            pos_threshold (float): Position threshold for convergence.
+            ori_threshold (float): Orientation threshold for convergence.
+            damping (float): Levenberg-Marquardt damping.
+            max_iters (int): Maximum number of iterations.
+            verbose (bool): Whether to print debug information.
+
+        Returns:
+            Tuple[bool, np.ndarray]: Success flag and the converged joint configuration.
+        """
+        if init_q is not None:
+            self._configuration.update(init_q)
+
+        end_effector_task = mink.FrameTask(
+            frame_name=site_name,
+            frame_type="site",
+            position_cost=1.0,
+            orientation_cost=1.0,
+            lm_damping=1.0,
+        )
+
+        end_effector_task.set_target(mink.SE3.from_matrix(target_pose))
+        tasks = [end_effector_task]
+
+        start_time = time.time()  # Start timing
+
+        for j in range(max_iters):
+            vel = mink.solve_ik(self._configuration, tasks, dt, solver, damping=damping, limits=limits)
+            self._configuration.integrate_inplace(vel, dt)
+            err = end_effector_task.compute_error(self._configuration)
+
+            pos_achieved = np.linalg.norm(err[:3]) <= pos_threshold
+            ori_achieved = np.linalg.norm(err[3:]) <= ori_threshold
+            if pos_achieved and ori_achieved:
+                end_time = time.time()  # End timing
+                elapsed_time = end_time - start_time
+                if verbose:
+                    print(
+                        f"Exiting after {j} iterations, configuration: {self._configuration.q}, time taken: {elapsed_time:.4f} seconds"
+                    )
+                return True, self._configuration.q
+
+        end_time = time.time()  # End timing
+        elapsed_time = end_time - start_time
+        if verbose:
+            print(
+                f"Failed to converge after {max_iters} iterations, time taken: {elapsed_time:.4f} seconds, pos_err: {err[:3]}, rot_err: {err[3:]}"
+            )
+        return False, self._configuration.q
+
+
+def main() -> None:
+    from i2rt.robots.motor_chain_robot import YAM_XML_PATH
+
+    mj_model = Kinematics(YAM_XML_PATH, "grasp_site")
+    q = np.zeros(6)
+    pose = mj_model.fk(q)
+    print(pose)
+
+    pose[0, 3] -= 0.1
+    pose[2, 3] += 0.1
+    print(pose)
+    q_ik = mj_model.ik(pose, "grasp_site")
+    print(q_ik)
+
+
+if __name__ == "__main__":
+    main()

i2rt/robots/motor_chain_robot.py

@@ -18,8 +18,8 @@
 from i2rt.utils.mujoco_utils import MuJoCoKDL
 
 I2RT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-YAM_XML_PATH = os.path.join(I2RT_ROOT, "robot_models/yam/yam.urdf")
-ARX_XML_PATH = os.path.join(I2RT_ROOT, "robot_models/arx_r5/arx.urdf")
+YAM_XML_PATH = os.path.join(I2RT_ROOT, "robot_models/yam/yam.xml")
+ARX_XML_PATH = os.path.join(I2RT_ROOT, "robot_models/arx_r5/arx.xml")
 
 
 @dataclass
@@ -358,6 +358,8 @@ def get_observations(self) -> Dict[str, np.ndarray]:
             if self._gripper_index is None:
                 return {
                     "joint_pos": self._joint_state.pos,
+                    "joint_vel": self._joint_state.vel,
+                    "joint_eff": self._joint_state.eff,
                 }
             else:
                 return {
@@ -416,7 +418,7 @@ def get_yam_robot(channel: str = "can0", model_path: str = YAM_XML_PATH) -> Moto
         motor_chain,
         xml_path=model_path,
         use_gravity_comp=True,
-        gravity_comp_factor=1.2,
+        gravity_comp_factor=1.3,
         gripper_index=6,
         gripper_limits=np.array([0.0, -2.7]),
         kp=np.array([80, 80, 80, 40, 10, 10, 20]),

i2rt/robots/tests/test_kinematics.py

@@ -0,0 +1,44 @@
+import numpy as np
+import pytest
+
+from i2rt.robots.kinematics import Kinematics
+from i2rt.robots.motor_chain_robot import YAM_XML_PATH
+
+
+@pytest.fixture
+def kinematics_yam() -> Kinematics:
+    return Kinematics(YAM_XML_PATH, "grasp_site")
+
+
+def test_fk(kinematics_yam: Kinematics) -> None:
+    q = np.zeros(6)
+    pose = kinematics_yam.fk(q)
+    assert pose.shape == (4, 4), "FK should return a 4x4 matrix"
+
+    # Add more assertions based on expected pose values
+    rotation = pose[:3, :3]
+    translation = pose[:3, 3]
+
+    start_rot = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
+    start_trans = np.array([0.245, 0.0, 0.164])
+    np.testing.assert_allclose(rotation, start_rot, atol=1e-5)
+    np.testing.assert_allclose(translation, start_trans, atol=1e-5)
+
+
+def test_ik_smoke(kinematics_yam: Kinematics) -> None:
+    q = np.zeros(6)
+    pose = kinematics_yam.fk(q)
+    pose[0, 3] -= 0.1
+    pose[2, 3] += 0.1
+    success, q_ik = kinematics_yam.ik(pose, "grasp_site")
+    assert success, "IK should succeed"
+    assert q_ik.shape == (6,), "IK should return a joint configuration of size 6"
+
+
+def test_cycle(kinematics_yam: Kinematics) -> None:
+    q = np.zeros(6)
+    pose = kinematics_yam.fk(q)
+    success, q_ik = kinematics_yam.ik(pose, "grasp_site")
+    assert success, "IK should succeed"
+    pose_reconstructed = kinematics_yam.fk(q_ik)
+    np.testing.assert_allclose(pose, pose_reconstructed, atol=1e-5)