Go2Py_SIM/Go2Py/robot/model.py

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_Jb_ = {}
        self.ef_Jw_ = {}

        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)
        self.M_ = None
        self.Minv_ = None
        self.nle_ = None
        self.g_ = None
        # 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 updateKinematics(self, q):
        """
        Updates the kinematic states. 

        Args:
            q (np.ndarray): A numpy array of size 19 representing the [x, y, z, qx, qy, qz, qw] and joint configurations in FR, FL, RR, RL order.
        """
        self.robot.computeJointJacobians(q)
        self.robot.framesForwardKinematics(q)\

        for ef_frame in self.ef_frames:
            Jw = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
            Jb = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL)
            self.ef_Jw_[ef_frame]=Jw[:, self.dq_reordering_idx]
            self.ef_Jb_[ef_frame]=Jb[:, self.dq_reordering_idx]
    
    def updateKinematicsPose(self, q, T):
        """
        Updates the kinematic states. 

        Args:
            q (np.ndarray): A numpy array of size 12 representing the joint configurations in FR, FL, RR, RL order.
            T (np.ndarray): 4x4 Transformation matrix representing the base pose of the robot.
        """
        q_ = np.hstack([pin.SE3ToXYZQUATtuple(pin.SE3(T)), q[self.q_reordering_idx]])
        self.robot.computeJointJacobians(q_)
        self.robot.framesForwardKinematics(q_)\

        for ef_frame in self.ef_frames:
            Jw = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
            Jb = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL)
            self.ef_Jw_[ef_frame]=Jw[:, self.dq_reordering_idx]
            self.ef_Jb_[ef_frame]=Jb[:, self.dq_reordering_idx]

    def updateDynamics(self, q, dq):
        """
        Updates the dynamical states. 

        Args:
            q (np.ndarray): A numpy array of size 19 representing the [x, y, z, qx, qy, qz, qw] and joint configurations in FR, FL, RR, RL order.
            dq (np.ndarray): A numpy array of size 18 representing the [vx, vy, vz, wx, wy, wz] and joint configurations in FR, FL, RR, RL order.
        """
        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]
        

    def updateAll(q, dq):
        """
        Updates the dynamic and kinematic parameters based on the given joint configurations and velocities.

        Args:
            q (np.ndarray): A numpy array of size 19 representing the [x, y, z, qx, qy, qz, qw] and joint configurations in FR, FL, RR, RL order.
            dq (np.ndarray): A numpy array of size 18 representing the [vx, vy, vz, wx, wy, wz] and joint configurations in FR, FL, RR, RL order.
        """
        self.updateKinematics(q)
        self.updateDynamics(q, dq)

    def updateAllPose(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.updateKinematics(q_)
        self.updateDynamics(q_, dq_)
        
    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_w':self.ef_Jw_,
            'J_b':self.ef_Jb_,
        }
    
    def getGroundReactionForce(self, tau_est, body_acceleration=None):
        if body_acceleration is None:   
            grf = {key:np.linalg.pinv(self.ef_J_[key][:3,6:].T)@(tau_est.squeeze() - self.nle_[6:]) for key in self.ef_J_.keys()}
        else:
            raise NotImplementedError("Ground reaction force with body dynamics is not implemented")
        return grf