[update] add joint smooth func to ik.

teleop/robot_control/joint_smooth.py

@@ -0,0 +1,92 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+class WeightedMovingFilter:
+    def __init__(self, weights, data_size = 14):
+        self._window_size = len(weights)
+        self._weights = np.array(weights)
+        assert np.isclose(np.sum(self._weights), 1.0), "[WeightedMovingFilter] the sum of weights list must be 1.0!"
+        self._data_size = data_size
+        self._filtered_data = np.zeros(self._data_size)
+        self._data_queue = []
+
+    def _apply_filter(self):
+        if len(self._data_queue) < self._window_size:
+            return self._data_queue[-1]
+
+        data_array = np.array(self._data_queue)
+        temp_filtered_data = np.zeros(self._data_size)
+        for i in range(self._data_size):
+            temp_filtered_data[i] = np.convolve(data_array[:, i], self._weights, mode='valid')[-1]
+        
+        return temp_filtered_data
+
+    def add_data(self, new_data):
+        assert len(new_data) == self._data_size
+        if len(self._data_queue) >= self._window_size:
+            self._data_queue.pop(0)
+
+        self._data_queue.append(new_data)
+        self._filtered_data = self._apply_filter()
+
+    @property
+    def filtered_data(self):
+        return self._filtered_data
+
+
+def visualize_filter_comparison(filter_params, steps):
+    import time
+    t = np.linspace(0, 4 * np.pi, steps)
+    original_data = np.array([np.sin(t + i) + np.random.normal(0, 0.2, len(t)) for i in range(35)]).T  # sin wave with noise, shape is [len(t), 35]
+
+    plt.figure(figsize=(14, 10))
+    
+    for idx, weights in enumerate(filter_params):
+        filter = WeightedMovingFilter(weights, 14)
+        data_2b_filtered = original_data.copy()
+        filtered_data = []
+
+        time1 = time.time()
+
+        for i in range(steps):
+            filter.add_data(data_2b_filtered[i][13:27])            # step i, columns 13 to 26 (total:14)
+            data_2b_filtered[i][13:27] = filter.filtered_data
+            filtered_data.append(data_2b_filtered[i])
+
+        time2 = time.time()
+        print(f"filter_params:{filter_params[idx]}, time cosume:{time2 - time1}")
+
+        filtered_data = np.array(filtered_data)
+
+        # col0 should not 2b filtered
+        plt.subplot(len(filter_params), 2, idx * 2 + 1)
+        plt.plot(filtered_data[:, 0], label=f'Filtered (Window {len(weights)})')
+        plt.plot(original_data[:, 0], 'r--', label='Original', alpha=0.5)
+        plt.title(f'Joint 1 - Should not to be filtered.')
+        plt.xlabel('Step')
+        plt.ylabel('Value')
+        plt.legend()
+
+        # col13 should 2b filtered
+        plt.subplot(len(filter_params), 2, idx * 2 + 2)
+        plt.plot(filtered_data[:, 13], label=f'Filtered (Window {len(weights)})')
+        plt.plot(original_data[:, 13], 'r--', label='Original', alpha=0.5)
+        plt.title(f'Joint 13 - Window {len(weights)}, Weights {weights}')
+        plt.xlabel('Step')
+        plt.ylabel('Value')
+        plt.legend()
+
+    plt.tight_layout()
+    plt.show()
+
+
+if __name__ == "__main__":
+    # windows_size and weights
+    filter_params = [ 
+        (np.array([0.7, 0.2, 0.1])),
+        (np.array([0.5, 0.3, 0.2])),
+        (np.array([0.4, 0.3, 0.2, 0.1])),
+    ]
+
+    visualize_filter_comparison(filter_params, steps = 100)

teleop/robot_control/robot_arm_ik.py

@@ -13,12 +13,14 @@ current_dir = os.path.dirname(os.path.abspath(__file__))
 parent_dir = os.path.dirname(current_dir)
 sys.path.append(parent_dir)
 
+from robot_control.joint_smooth import WeightedMovingFilter
+
 class Arm_IK:
     def __init__(self):
         np.set_printoptions(precision=5, suppress=True, linewidth=200)
 
-        # self.robot = pin.RobotWrapper.BuildFromURDF('../assets/g1/g1_body29_hand14.urdf', '../assets/g1/')
-        self.robot = pin.RobotWrapper.BuildFromURDF('../../assets/g1/g1_body29_hand14.urdf', '../../assets/g1/') # for test
+        self.robot = pin.RobotWrapper.BuildFromURDF('../assets/g1/g1_body29_hand14.urdf', '../assets/g1/')
+        # self.robot = pin.RobotWrapper.BuildFromURDF('../../assets/g1/g1_body29_hand14.urdf', '../../assets/g1/') # for test
 
         self.mixed_jointsToLockIDs = [  
                                         "left_hip_pitch_joint" ,
@@ -74,7 +76,8 @@ class Arm_IK:
                               np.array([0.05,0,0]).T),
                       pin.FrameType.OP_FRAME)
         )
-        
+
+        self.joint_smooth_filter = WeightedMovingFilter(np.array([0.4, 0.3, 0.2, 0.1]), 14)
 
         self.init_data = np.zeros(self.reduced_robot.model.nq)
 
@@ -202,6 +205,9 @@ class Arm_IK:
             sol = self.opti.solve_limited()
             sol_q = self.opti.value(self.var_q)
 
+            self.joint_smooth_filter.add_data(sol_q)
+            sol_q = self.joint_smooth_filter.filtered_data
+
             # self.vis.display(sol_q)  # for visualization
             self.init_data = sol_q