diff --git a/Go2Py/robot/model.py b/Go2Py/robot/model.py index a534c3b..196f2ae 100644 --- a/Go2Py/robot/model.py +++ b/Go2Py/robot/model.py @@ -1 +1,175 @@ +import os +from Go2Py import ASSETS_PATH import pinocchio as pin +import numpy as np +urdf_path = os.path.join(ASSETS_PATH, 'urdf/go2.urdf') +urdf_root_path = os.path.join(ASSETS_PATH, 'urdf') + + +class Go2Model: + """ + A model class for the Go2 quadruped robot using the Pinocchio library. + + Attributes: + robot (pin.RobotWrapper): The Pinocchio RobotWrapper instance for the Go2 robot. + data (pin.Model.Data): The data structure used by Pinocchio for computations. + ef_frames (list): List of end-effector frame names. + dq_reordering_idx (np.ndarray): Index array for reordering the joint velocity vector. + q_reordering_idx (np.ndarray): Index array for reordering the joint position vector. + ef_J_ (dict): A dictionary storing the Jacobians for the end-effector frames. + """ + def __init__(self): + """ + Initializes the Go2Model class by loading the URDF, setting up the robot model, and calculating initial dimensions. + """ + self.robot = pin.RobotWrapper.BuildFromURDF(urdf_path, urdf_root_path, pin.JointModelFreeFlyer()) + self.data = self.robot.data + # Standing joint configuration in Unitree Joint order + self.ef_frames = ['FR_foot', 'FL_foot', 'RR_foot', 'RL_foot'] + self.dq_reordering_idx = np.array([0, 1, 2, 3, 4, 5,\ + 9, 10, 11, 6, 7, 8, 15, 16, 17, 12, 13, 14]) + self.q_reordering_idx = np.array([9, 10, 11, 6, 7, 8, 15, 16, 17, 12, 13, 14])-6 + self.ef_J_ = {} + + ID_FL_HAA = self.robot.model.getFrameId('FL_hip_joint') + ID_FR_HAA = self.robot.model.getFrameId('FR_hip_joint') + ID_RL_HAA = self.robot.model.getFrameId('RL_hip_joint') + ID_RR_HAA = self.robot.model.getFrameId('RR_hip_joint') + ID_FL_HFE = self.robot.model.getFrameId('FL_thigh_joint') + ID_FR_HFE = self.robot.model.getFrameId('FR_thigh_joint') + ID_RL_HFE = self.robot.model.getFrameId('RL_thigh_joint') + ID_RR_HFE = self.robot.model.getFrameId('RR_thigh_joint') + ID_FL_KFE = self.robot.model.getFrameId('FL_calf_joint') + ID_FR_KFE = self.robot.model.getFrameId('FR_calf_joint') + ID_RL_KFE = self.robot.model.getFrameId('RL_calf_joint') + ID_RR_KFE = self.robot.model.getFrameId('RR_calf_joint') + ID_FR_FOOT = self.robot.model.getFrameId('FR_foot') + + q_neutral = np.asarray([0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) + self.robot.framesForwardKinematics(q_neutral) + + self.h = np.linalg.norm(self.robot.data.oMf[ID_FR_HAA].translation - self.robot.data.oMf[ID_RR_HAA].translation) + self.b = np.linalg.norm(self.robot.data.oMf[ID_FR_HAA].translation - self.robot.data.oMf[ID_FL_HAA].translation) + self.l1 = np.linalg.norm(self.robot.data.oMf[ID_FR_HAA].translation - self.robot.data.oMf[ID_FR_HFE].translation) + self.l2 = np.linalg.norm(self.robot.data.oMf[ID_FR_HFE].translation - self.robot.data.oMf[ID_FR_KFE].translation) + self.l3 = np.linalg.norm(self.robot.data.oMf[ID_FR_KFE].translation - self.robot.data.oMf[ID_FR_FOOT].translation) + + print(self.robot.data.oMf[ID_FR_HAA].translation - self.robot.data.oMf[ID_RR_HAA].translation) + + print(self.h) + print(self.b) + print(self.l1) + print(self.l2) + print(self.l3) + + def inverseKinematics(self, T, feet_pos): + """ + Calculates the inverse kinematics for the robot given a desired state. + + Args: + T (np.ndarray): The 4x4 homogenous transformation representing the pose of the base_link in the world frame + x (np.ndarray): A numpy array of size 12 representing foot positions in world frame in FR, FL, RR, RL order. + + Returns: + np.ndarray: A numpy array of size 12 representing the joint angles of the legs. + """ + rB = np.asarray(T[0:3, -1]) # Base position (3D vector) + R = T[:3,:3] # Body orientation (quaternion converted to rotation matrix) + + sx = [1, 1, -1, -1] + sy = [-1, 1, -1, 1] + + joint_angles = np.zeros(12) + + for i in range(4): + r_HB = np.array([sx[i] * self.h / 2, sy[i] * self.b / 2, 0]) # Hip offset (3D vector) + rf = np.asarray(feet_pos[3*i:3*i+3]) # Foot position (3D vector) + r_fH = R.T @ (rf - rB) - r_HB # Foot relative to hip in body frame (3D vector) + + x = r_fH[0] + y = r_fH[1] + z = r_fH[2] + et = y**2 + z**2 - self.l1**2 + + # Theta 3 calculation + c3 = (x**2 + et - self.l2**2 - self.l3**2) / (2 * self.l2 * self.l3) + s3 = -np.sqrt(1 - c3**2) + t3 = np.arctan2(s3, c3) + + # Theta 2 calculation + k1 = self.l2 + self.l3 * c3 + k2 = self.l3 * s3 + r1 = np.sqrt(k1**2 + k2**2) + t2 = np.arctan2(-x/r1, np.sqrt(et) / r1) - np.arctan2(k2/r1, k1/r1) + + # Theta 1 calculation + zv = self.l2 * np.cos(t2) + self.l3 * np.cos(t2 + t3) + m1 = sy[i] * self.l1 + m2 = -zv + r2 = np.sqrt(m1**2 + m2**2) + t1 = np.arctan2(z/r2, y/r2) - np.arctan2(m2/r2, m1/r2) + + joint_angles[3*i:3*i+3] = np.array([t1, t2, t3]) + + # TODO: Implement joint axis direction multiplication from URDF + + return joint_angles + + def forwardKinematics(self, T, q): + """ + Computes the forward kinematics for the robot given a transformation matrix and joint configuration. + + Args: + T (np.ndarray): 4x4 Transformation matrix representing the base pose of the robot. + q (np.ndarray): A numpy array of size 12 representing the joint configurations in FR, FL, RR, RL order. + + Returns: + dict: A dictionary containing the poses of specified frames in the robot. + """ + q_ = np.hstack([pin.SE3ToXYZQUATtuple(pin.SE3(T)), q[self.q_reordering_idx]]) + self.robot.framesForwardKinematics(q_) + ef_frames = ['base_link','FR_foot', 'FL_foot', 'RR_foot', 'RL_foot'] + data = {} + return {frame:self.robot.data.oMf[self.robot.model.getFrameId(frame)].homogeneous \ + for frame in ef_frames} + + + def update(self, q, dq, T, v): + """ + Updates the dynamic and kinematic parameters based on the given joint configurations and velocities. + + Args: + q (np.ndarray): A numpy array of size 12 representing the joint configurations in FR, FL, RR, RL order. + dq (np.ndarray): A numpy array of size 12 representing the joint velocities in FR, FL, RR, RL order. + T (np.ndarray): 4x4 Transformation matrix representing the base pose of the robot. + v (np.ndarray): A numpy array of size 6 representing the base velocity in body frame [v, w]. + """ + q_ = np.hstack([pin.SE3ToXYZQUATtuple(pin.SE3(T)), q[self.q_reordering_idx]]) + dq_ = np.hstack([v, dq[self.q_reordering_idx]]) + self.robot.computeJointJacobians(q_) + self.robot.framesForwardKinematics(q_) + self.robot.centroidalMomentum(q_,dq_) + self.nle_ = self.robot.nle(q_, dq_)[self.dq_reordering_idx] + self.g_ = self.robot.gravity(q_)[self.dq_reordering_idx] + self.M_ = self.robot.mass(q_)[self.dq_reordering_idx,:] + self.M_ = self.M_[:,self.dq_reordering_idx] + self.Minv_ = pin.computeMinverse(self.robot.model, self.robot.data, q_)[self.dq_reordering_idx,:] + self.Minv_ = self.Minv_[:,self.dq_reordering_idx] + for ef_frame in self.ef_frames: + J = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL) + self.ef_J_[ef_frame]=J[:, self.dq_reordering_idx] + + def getInfo(self): + """ + Retrieves the current dynamics and kinematic information of the robot. + + Returns: + dict: A dictionary containing the robot's mass matrix, inverse mass matrix, non-linear effects, gravity vector, and Jacobians for the end-effectors. + """ + return { + 'M':self.M_, + 'Minv':self.Minv_, + 'nle':self.nle_, + 'g':self.g_, + 'J':self.ef_J_, + } \ No newline at end of file