MassageRobot_Dobot/Massage/MassageControl/algorithms/hybridPid_controller.py

import sys
import numpy as np
from scipy.spatial.transform import Rotation as R
from .arm_state import ArmState
from .base_controller import BaseController
from pathlib import Path
sys.path.append(str(Path(__file__).resolve().parent.parent))
from tools.yaml_operator import read_yaml
import time

class HybridPidController(BaseController):
    def __init__(self, name, state: ArmState,config_path) -> None:
        super().__init__(name, state)
        self.load_config(config_path)
        self.laset_print_time = 0
    

    def load_config(self, config_path):
        config_dict = read_yaml(config_path)
        if self.name != config_dict['name']:
            raise ValueError(f"Controller name {self.name} does not match config name {config_dict['name']}")
        # 姿态调节器
        # 位控 2x2矩阵
        self.Kp_R = np.diag(np.array(config_dict['Kp_R']))
        self.Ki_R = np.diag(np.array(config_dict['Ki_R']))
        self.Kd_R = np.diag(np.array(config_dict['Kd_R']))

        # 力控
        self.force_mass = np.array(config_dict['force_mass'])
        self.force_damp = np.array(config_dict['force_damp'])
        self.Kp_force = np.array(config_dict['Kp_force'])
        self.Kd_force = np.array(config_dict['Kd_force'])
        self.Ki_force = np.array(config_dict['Ki_force'])
        # TODO 修改控制率
        self.e_t = 0
        self.e_t1 = 0
        self.e_t2 = 0
        self.force_control_value = 0

        # 位控 2x2矩阵
        self.Kp = np.diag(np.array(config_dict['Kp']))
        self.Ki = np.diag(np.array(config_dict['Ki']))
        self.Kd = np.diag(np.array(config_dict['Kd']))
        self.pose_integral_error = np.zeros(6)
    
    def step(self,dt):
        # 方向统一
        if self.state.desired_orientation.dot(self.state.arm_orientation) < 0:
            self.state.arm_orientation = -self.state.arm_orientation
        # 缓存常用计算
        arm_ori_quat = R.from_quat(self.state.arm_orientation)
        arm_ori_mat = arm_ori_quat.as_matrix()
        # 位置误差
        temp_pose_error = self.state.arm_position - self.state.desired_position
        self.state.pose_error[:3] = arm_ori_mat.T @ temp_pose_error
        # 姿态误差(四元数)
        rot_err_quat = arm_ori_quat.inv() * R.from_quat(self.state.desired_orientation)
        self.state.pose_error[3:] = -rot_err_quat.as_rotvec(degrees=False)
        # 期望加速度
        wrench_err = self.state.external_wrench_tcp - self.state.desired_wrench

        # z方向力控制
        self.e_t2 = self.e_t1
        self.e_t1 = self.e_t
        self.e_t = wrench_err[2]
        self.force_control_value += self.Kp_force * (self.e_t - self.e_t1) + self.Ki_force * self.e_t * dt + self.Kd_force * (self.e_t - 2 * self.e_t1 + self.e_t2) /dt
        self.state.arm_desired_acc[2] = (1.0 / self.force_mass) * (self.force_control_value - self.force_damp * self.state.arm_desired_twist[2])

        self.pose_integral_error +=  self.state.pose_error * dt
        # 位控制
        self.state.arm_desired_acc[:2] = -self.Kd @ self.state.arm_desired_twist[:2] - self.Kp @ self.state.pose_error[:2] - self.Ki @ self.pose_integral_error[:2]
        # 姿态pid
        self.state.arm_desired_acc[3:] = -self.Kd_R @ self.state.arm_desired_twist[3:] - self.Kp_R @ self.state.pose_error[3:] - self.Ki_R @ self.pose_integral_error[3:]
        self.clip_command(self.state.arm_desired_acc, "acc")
        ## 更新速度和位姿
        self.state.arm_desired_twist += self.state.arm_desired_acc * dt
        self.clip_command(self.state.arm_desired_twist, "vel")
        # 计算位姿变化
        delta_pose = np.concatenate([
            arm_ori_mat @ (self.state.arm_desired_twist[:3] * dt),
            self.state.arm_desired_twist[3:] * dt
        ])
        self.clip_command(delta_pose, "pose")
        # 更新四元数
        delta_ori_quat = R.from_rotvec(delta_pose[3:]).as_quat()
        arm_ori_quat_new = arm_ori_quat * R.from_quat(delta_ori_quat)
        self.state.arm_orientation_command = arm_ori_quat_new.as_quat()
        # 归一化四元数
        self.state.arm_orientation_command /= np.linalg.norm(self.state.arm_orientation_command)
        # 更新位置
        self.state.arm_position_command = self.state.arm_position + delta_pose[:3]

    def step_traj(self,dt):
        # 方向统一
        if self.state.desired_orientation.dot(self.state.arm_orientation) < 0:
            self.state.arm_orientation = -self.state.arm_orientation
        # 缓存常用计算
        arm_ori_quat = R.from_quat(self.state.arm_orientation)
        arm_ori_mat = arm_ori_quat.as_matrix()
        # 位置误差
        
        temp_pose_error = self.state.arm_position - self.state.desired_position + self.state.desired_twist[:3] * dt
        self.state.pose_error[:3] = arm_ori_mat.T @ temp_pose_error

        # 姿态误差(四元数)
        angular_velocity = np.array(self.state.desired_twist[3:])  # 形状 (3,)

        # 用旋转向量（小角度近似）
        rotvec = angular_velocity * dt  # 旋转向量 = 角速度 × 时间
        rot_quat = R.from_rotvec(rotvec).as_quat()  # 转成四元数，形状 (4,)
        rot_err_quat = R.from_quat(rot_quat).inv() * arm_ori_quat.inv() * R.from_quat(self.state.desired_orientation)
        self.state.pose_error[3:] = -rot_err_quat.as_rotvec(degrees=False)

        # 期望加速度
        wrench_err = self.state.external_wrench_tcp - self.state.desired_wrench

        # z方向力控制
        self.e_t2 = self.e_t1
        self.e_t1 = self.e_t
        self.e_t = wrench_err[2]
        self.force_control_value += self.Kp_force * (self.e_t - self.e_t1) + self.Ki_force * self.e_t * dt + self.Kd_force * (self.e_t - 2 * self.e_t1 + self.e_t2) /dt
        self.state.arm_desired_acc[2] = (1.0 / self.force_mass) * (self.force_control_value - self.force_damp * self.state.arm_desired_twist[2])
        self.pose_integral_error +=  self.state.pose_error * dt
        # 位控制
        self.state.arm_desired_acc[:2] = -self.Kd @ self.state.arm_desired_twist[:2] - self.Kp @ self.state.pose_error[:2] - self.Ki @ self.pose_integral_error[:2]
        # 姿态pid
        self.state.arm_desired_acc[3:] = -self.Kd_R @ self.state.arm_desired_twist[3:] - self.Kp_R @ self.state.pose_error[3:] - self.Ki_R @ self.pose_integral_error[3:]

        self.clip_command(self.state.arm_desired_acc, "acc")
        ## 更新速度和位姿
        self.state.arm_desired_twist += self.state.arm_desired_acc * dt
        self.clip_command(self.state.arm_desired_twist, "vel")
        # 计算位姿变化
        delta_pose = np.concatenate([
            arm_ori_mat @ (self.state.arm_desired_twist[:3] * dt),
            self.state.arm_desired_twist[3:] * dt
        ])
        self.clip_command(delta_pose, "pose")
        # 更新四元数
        delta_ori_quat = R.from_rotvec(delta_pose[3:]).as_quat()
        arm_ori_quat_new = arm_ori_quat * R.from_quat(delta_ori_quat)
        self.state.arm_orientation_command = arm_ori_quat_new.as_quat()
        # 归一化四元数
        self.state.arm_orientation_command /= np.linalg.norm(self.state.arm_orientation_command)
        # 更新位置
        self.state.arm_position_command = self.state.arm_position + delta_pose[:3]