From a91962aa3a870608c1397b7bc119007a24deedc3 Mon Sep 17 00:00:00 2001
From: Iori Yanokura <ab.ioryz@gmail.com>
Date: Wed, 10 Apr 2024 22:21:33 +0900
Subject: [PATCH] Add trim vector to adjust real robot angles

---
 rcb4/armh7interface.py                        | 54 ++++++++++++++++++-
 rcb4/rcb4interface.py                         |  6 ++-
 .../node_scripts/rcb4_ros_bridge.py           |  8 ++-
 .../config/kxrmw4a6h2m_servo_config.yaml      |  2 +-
 tests/rcb4_tests/test_armh7interface.py       |  7 ++-
 tests/rcb4_tests/test_rcb4interface.py        |  4 +-
 6 files changed, 72 insertions(+), 9 deletions(-)

diff --git a/rcb4/armh7interface.py b/rcb4/armh7interface.py
index 9b9017a7..839e5845 100644
--- a/rcb4/armh7interface.py
+++ b/rcb4/armh7interface.py
@@ -24,6 +24,7 @@
 from rcb4.ctype_utils import c_type_to_size
 from rcb4.data import kondoh7_elf
 from rcb4.rcb4interface import CommandTypes
+from rcb4.rcb4interface import deg_to_servovector
 from rcb4.rcb4interface import rcb4_dof
 from rcb4.rcb4interface import RCB4Interface
 from rcb4.rcb4interface import ServoParams
@@ -516,6 +517,46 @@ def angle_vector_to_servo_angle_vector(self, av, servo_ids=None):
         tmp_av[seq_indices] = np.array(av)
         return np.matmul(self.joint_to_actuator_matrix, tmp_av)[seq_indices]
 
+    def _trim_servo_vector(self):
+        return np.array(
+            self.read_cstruct_slot_vector(ServoStruct, 'trim'),
+            dtype=np.int16)
+
+    def _set_trim_vector(self, trim_vector=None, servo_ids=None):
+        if servo_ids is None:
+            servo_ids = self.search_servo_ids()
+        if len(trim_vector) != len(servo_ids):
+            raise ValueError(
+                'Length of servo_ids and trim_vector must be the same.')
+        if len(trim_vector) == 0:
+            return
+        current_trim_vector = self._trim_servo_vector()
+        current_trim_vector[np.array(servo_ids)] = trim_vector
+        self.write_cstruct_slot_v(ServoStruct, 'trim',
+                                  current_trim_vector)
+
+    def trim_vector(self, av=None, servo_ids=None):
+        if av is not None:
+            trim_av = self.angle_vector_to_servo_angle_vector(
+                av, servo_ids) - 7500
+            return self._set_trim_vector(trim_av, servo_ids)
+        all_servo_ids = self.search_servo_ids()
+        if len(all_servo_ids) == 0:
+            return np.empty(shape=0)
+        trim = np.append(7500 + self._trim_servo_vector()[all_servo_ids], 1)
+        trim_av = np.matmul(trim.T, self.actuator_to_joint_matrix.T)[:-1]
+        if servo_ids is not None:
+            if len(servo_ids) == 0:
+                return np.empty(shape=0)
+            trim_av = trim_av[self.sequentialized_servo_ids(servo_ids)]
+        return trim_av
+
+    def clear_trim_vector(self, write_to_flash=True):
+        self.write_cstruct_slot_v(
+            ServoStruct, 'trim', np.zeros(rcb4_dof))
+        if write_to_flash:
+            self.write_to_flash()
+
     def ics_start(self):
         return self.set_cstruct_slot(SystemStruct, 0, 'ics_comm_stop',
                                      [0, 0, 0, 0, 0, 0])
@@ -874,7 +915,13 @@ def buzzer(self):
         return self.cfunc_call('buzzer_init_sound', [])
 
     def write_to_flash(self):
-        return self.cfunc_call('rom_to_flash', [])
+        # The operation write_to_flash is time-consuming,
+        # hence the default_timeout is temporarily extended.
+        default_timeout = self._default_timeout
+        self._default_timeout = 5.0
+        ret = self.cfunc_call('rom_to_flash', [])
+        self._default_timeout = default_timeout
+        return ret
 
     def set_data_address(self):
         self.set_sidata()
@@ -970,6 +1017,9 @@ def write_cstruct_slot_v(self, cls, slot_name, vec):
                 elif typ == 'uint16':
                     v = int(round(v))
                     struct.pack_into('<H', byte_list, 10 + i * esize, v)
+                elif typ == 'int16':
+                    v = int(round(v))
+                    struct.pack_into('<h', byte_list, 10 + i * esize, v)
                 elif typ == 'uint32':
                     v = int(round(v))
                     struct.pack_into('<I', byte_list, 10 + i * esize, v)
@@ -1023,7 +1073,7 @@ def joint_to_actuator_matrix(self):
             self._joint_to_actuator_matrix[:, servo_length] = 7500
             self._joint_to_actuator_matrix[servo_length, servo_length] = 1
             for i in range(servo_length):
-                self._joint_to_actuator_matrix[i, i] = 30
+                self._joint_to_actuator_matrix[i, i] = deg_to_servovector
         return self._joint_to_actuator_matrix
 
     @property
diff --git a/rcb4/rcb4interface.py b/rcb4/rcb4interface.py
index 3bbe6793..7158cbe6 100644
--- a/rcb4/rcb4interface.py
+++ b/rcb4/rcb4interface.py
@@ -14,6 +14,10 @@
 from rcb4.asm import rcb4_velocity
 
 
+# 4000.0 / 135
+deg_to_servovector = 29.62962962962963
+
+
 class CommandTypes(Enum):
     Move = 0x00
     Jump = 0x0B
@@ -600,7 +604,7 @@ def joint_to_actuator_matrix(self):
             self._joint_to_actuator_matrix[:, servo_length] = 7500
             self._joint_to_actuator_matrix[servo_length, servo_length] = 1
             for i in range(servo_length):
-                self._joint_to_actuator_matrix[i, i] = 30
+                self._joint_to_actuator_matrix[i, i] = deg_to_servovector
         return self._joint_to_actuator_matrix
 
     @property
diff --git a/ros/kxr_controller/node_scripts/rcb4_ros_bridge.py b/ros/kxr_controller/node_scripts/rcb4_ros_bridge.py
index 8096b728..7702e48c 100644
--- a/ros/kxr_controller/node_scripts/rcb4_ros_bridge.py
+++ b/ros/kxr_controller/node_scripts/rcb4_ros_bridge.py
@@ -152,11 +152,14 @@ def __init__(self):
                         self.interface.search_servo_ids().tolist())
 
         wheel_servo_sorted_ids = []
+        trim_vector_servo_ids = []
+        trim_vector_offset = []
         for _, info in servo_infos.items():
             if isinstance(info, int):
                 continue
             servo_id = info['id']
-            direction = info['direction']
+            direction = info.get('direction', 1)
+            offset = info.get('offset', 0)
             if 'type' in info and info['type'] == 'continuous':
                 wheel_servo_sorted_ids.append(servo_id)
             idx = self.interface.servo_id_to_index(servo_id)
@@ -164,6 +167,9 @@ def __init__(self):
                 continue
             self.interface._joint_to_actuator_matrix[idx, idx] = \
                 direction * self.interface._joint_to_actuator_matrix[idx, idx]
+            trim_vector_servo_ids.append(servo_id)
+            trim_vector_offset.append(direction * offset)
+        self.interface.trim_vector(trim_vector_offset, trim_vector_servo_ids)
         if self.interface.wheel_servo_sorted_ids is None:
             self.interface.wheel_servo_sorted_ids = wheel_servo_sorted_ids
 
diff --git a/ros/kxr_models/config/kxrmw4a6h2m_servo_config.yaml b/ros/kxr_models/config/kxrmw4a6h2m_servo_config.yaml
index 4eb8dcef..834ff58a 100644
--- a/ros/kxr_models/config/kxrmw4a6h2m_servo_config.yaml
+++ b/ros/kxr_models/config/kxrmw4a6h2m_servo_config.yaml
@@ -1,5 +1,5 @@
 joint_name_to_servo_id: {
-  HEAD_JOINT0: 32,
+  HEAD_JOINT0: {id: 32, direction: 1, offset: 0},
   HEAD_JOINT1: 34,
   LARM_JOINT0: 3,
   LARM_JOINT1: 5,
diff --git a/tests/rcb4_tests/test_armh7interface.py b/tests/rcb4_tests/test_armh7interface.py
index aa8f17c3..917b45c9 100644
--- a/tests/rcb4_tests/test_armh7interface.py
+++ b/tests/rcb4_tests/test_armh7interface.py
@@ -41,12 +41,12 @@ def test_angle_vector_to_servo_angle_vector(self):
         sv = self.interface.angle_vector_to_servo_angle_vector(
             [30, 60], [32, 34])
         testing.assert_array_almost_equal(
-            sv, [8400, 9300])
+            sv, [8388.888874, 9277.777748])
 
         sv = self.interface.angle_vector_to_servo_angle_vector(
             [30, 60], [34, 32])
         testing.assert_array_almost_equal(
-            sv, [8400, 9300])
+            sv, [8388.888874, 9277.777748])
 
     def test_angle_vector(self):
         self.interface.hold()
@@ -105,3 +105,6 @@ def test_servo_angle_vector(self):
         self.interface.neutral()
         time.sleep(4.0)
         self.interface.free()
+
+    def test_trim_vector(self):
+        self.interface.trim_vector()
diff --git a/tests/rcb4_tests/test_rcb4interface.py b/tests/rcb4_tests/test_rcb4interface.py
index da111aec..88184233 100644
--- a/tests/rcb4_tests/test_rcb4interface.py
+++ b/tests/rcb4_tests/test_rcb4interface.py
@@ -41,12 +41,12 @@ def test_angle_vector_to_servo_angle_vector(self):
         sv = self.interface.angle_vector_to_servo_angle_vector(
             [30, 60], [32, 34])
         testing.assert_array_almost_equal(
-            sv, [8400, 9300])
+            sv, [8388.888874, 9277.777748])
 
         sv = self.interface.angle_vector_to_servo_angle_vector(
             [30, 60], [34, 32])
         testing.assert_array_almost_equal(
-            sv, [8400, 9300])
+            sv, [8388.888874, 9277.777748])
 
     def test_angle_vector(self):
         self.interface.hold()