A simple Quadrotor simulation with a PD controller

quadrotor_dynamics.py

"""
A simple quadrotor simulatior. The code should self-explanatory.

Reference: Quadcopter Dynamics, Simulation, and Control. by Andrew Gibiansky.

Download: chrome-extension://oemmndcbldboiebfnladdacbdfmadadm/http://andrew.gibiansky.com/downloads/pdf/Quadcopter%20Dynamics,%20Simulation,%20and%20Control.pdf

"""


import numpy as np
#
R_x = lambda phi : np.array([[1., 0., 0.],
                             [0., np.cos(phi), -np.sin(phi)],
                             [0., np.sin(phi), np.cos(phi)]])
#
R_y = lambda theta : np.array([[np.cos(theta), 0., np.sin(theta)],
                               [0., 1., 0.],
                               [-np.sin(theta), 0, np.cos(theta)]])
#
R_z = lambda psi : np.array([[np.cos(psi), -np.sin(psi), 0.],
                             [np.sin(psi), np.cos(psi), 0.],
                             [0., 0., 1.]])
#
class PDCtrl(object):

    def __init__(self, kd=4., kp=3.):
        self.kd = kd
        self.kp = kp

    def __call__(self, theta_dot, theta, drone):
        """
        To control the angular velocities and the Euler angles
        """
        # compute total thrust
        T = drone.m * drone.g / (drone.k * np.cos(theta[0, 0]) * np.cos(theta[1, 0]))
        #
        error = self.kd * (theta_dot - 0) + self.kp * theta
        #
        e1, e2, e3 = error[0, 0], error[1, 0], error[2, 0]
        k, L, b = drone.k, drone.L, drone.b
        Ixx, Iyy, Izz = drone.I[0, 0], drone.I[1, 1], drone.I[2, 2]
        actions = np.zeros(shape=(4, 1))
        actions[0, 0] = T/4 - (2*b*e1*Ixx+e3*Izz*k*L) / (4*b*k*L)
        actions[1, 0] = T/4 + (e3*Izz)/(4*b) - (e2*Iyy)/(2*k*L)
        actions[2, 0] = T/4 - (-2*b*e1*Ixx+e3*Izz*k*L)/(4*b*k*L)
        actions[3, 0] = T/4 + (e3*Izz)/(4*b) + (e2*Iyy)/(2*k*L)
        #
        return actions


class QuadrotorDynamics(object):

    def __init__(self, ):
        # Some basic Parameters
        self.g = 9.81 # gravity acceleration
        self.m = 0.5  # mass
        self.L = 0.25 # length of
        self.k = 3e-6 # thrust coefficient
        self.kd = 0.25 # drag coefficient
        self.b = 1e-7
        # #
        self.I = np.diag([5e-3, 5e-3, 10e-3]) # inertial matrix
        #
        self.t0, self.tend = 0., 10
        self.dt = 0.005

    def run_simulation(self, controller=None):
        #
        ts = np.arange(self.t0, self.tend, self.dt)
        N = ts.shape[0]
        #
        pos = np.array([1, 1, 1])[:, np.newaxis]  # Position of the drone
        pos_dot = np.zeros(shape=(3, 1))          # First derivation, aka velocity
        theta = np.ones(shape=(3, 1))            # Eular Angles that represents the orientation in inertial frame
        theta_dot = np.ones(shape=(3, 1))        # First derivation, aka Angular velocity
        #
        orientation = np.zeros([N, 6])
        for i, t_i in enumerate(ts):
            # # # # # # # # # #
            # get actions
            # # # # # # # # # #
            # the angular velocities of the rotors
            # actions = self.get_actions()
            actions = controller(theta_dot, theta, self)

            # convert angular velocities from inertial frame to body frame
            omega = self.angles_dot_convert(theta, theta_dot, convert='w_to_b');
            # compute the acceleration of the rigidy body
            acc = self.get_accelarations(actions, omega, pos_dot)
            # compute the angular velocities of the rigid body
            omega_dot = self.get_angular_accelerations(actions, omega)

            # integrate angular velocities
            omega = omega + self.dt * omega_dot
            # convert angular velocities from body frame to inertial frame
            theta_dot = self.angles_dot_convert(theta, omega, convert="b_to_w")
            # integrate Eular angulers, which represents the orientation of rigid body
            theta = theta + self.dt * theta_dot
            # integrate linear accelerations in order to get the linear velocity
            pos_dot = pos_dot + self.dt * acc
            # integrate linear velocity in order to get the linear position
            pos = pos + self.dt * pos_dot

            #
            theta_x, theta_y, theta_z = theta[:, 0]
            theta_dotx, theta_doty, theta_dotz = theta_dot[:, 0]
            orientation[i, :3] = (theta_x, theta_y, theta_z)
            orientation[i, 3:6] = (theta_dotx, theta_doty, theta_dotz)
        return orientation


    @staticmethod
    def get_actions():
        """
        Actions, which are propotional to the angular velocity of propellers
        which are vectors pointing along the axis of rotation.

            u = action ** 2

        Angular velocity is different from the time derivative of roll,
        pitch, and roll, which is also termed angular velocities in some cases.
        params:
        ----------
            t : dummy params, not used.
        """
        action = np.ones(shape=(4, 1)) * 640
        action[0, 0] += 100
        return action**2

    def get_thrust(self, actions):
        """
        Compute thrust on the quadrotor (in the body frame)
        params
        ----------
            action : current action, or the angular velocity
        """
        return np.array([0., 0., self.k * np.sum(actions)])[:, np.newaxis]

    def get_torques(self, actions):
        """
        Compute torques, given current angular velocity of rotors
        """
        tau = np.array([[self.L * self.k * (actions[0, 0] - actions[2, 0])],
                        [self.L * self.k * (actions[1, 0] - actions[3, 0])],
                        [self.b * (actions[0,0] - actions[1,0] + actions[2,0] - actions[3,0])],
                       ])
        return tau

    def get_angular_accelerations(self, actions, theta_dot):
        """
        Compute angular accelerations in body frame.
        params
        ---------
        """
        # compute torques
        tau = self.get_torques(actions)
        theta_ddot = np.linalg.inv(self.I) @ (tau - np.cross(theta_dot[:,0], self.I@theta_dot[:,0])[:, np.newaxis])
        return theta_ddot

    def get_accelarations(self, actions, theta, state_dot):
        """
        Compute linear accelerations
        params
        ----------
            actions : current actions, angular velocity of the four rotors
            theta : angles in the inertial frame
            state_dot : linear velocities of the body
        """
        gravity = np.array([0., 0., -self.g])[:, np.newaxis]
        # extract roll, pitch, and yaw
        phi_x, theta_y, psi_z = theta[:, 0]
        # compute the rotation matrix for get_thrust
        R_xyz = R_z(psi_z) @ R_y(theta_y) @ R_x(phi_x)
        # compute torques on the quadrotor
        TB = self.get_thrust(actions)
        # compute the drag forces
        F_D = - self.kd * state_dot
        return gravity + R_xyz @ TB / self.m + F_D

    @staticmethod
    def angles_dot_convert(theta, theta_dot, convert="w_to_b"):
        """
        Transformation of the angular velocities from the inertial frame {w}
        to the body frame {b}.
            \dot{\Theta}_B = W_{\Theta_{BW}} * \dot{\Theta}_W
        params
        ----------
            theta     : body orientation in inertial frame, (roll, pitch, yaw)
            theta_dot : angular velocities
            convert    : either from inertial to body, or from body to inertial frame
        """
        # extract roll, pitch, and yaw
        phi_x, theta_y, psi_z = theta[:, 0]
        # Transformation matrix for thetaular velocities
        # from iniertial frame to body frame
        W_bw = np.array([[1, 0, -np.sin(theta_y)],
                         [0, np.cos(phi_x), np.cos(theta_y)*np.sin(phi_x)],
                         [0, -np.sin(phi_x), np.cos(theta_y)*np.cos(phi_x)]
                       ])
        # compute the angular velocities in body frame
        if convert == "w_to_b":
            return W_bw @ theta_dot
        elif convert == "b_to_w":
            return np.linalg.inv(W_bw) @ theta_dot
        else:
            raise NotImplementedError


def main():
    #
    pd_controller = PDCtrl()
    #
    quadrotor_dynamics = QuadrotorDynamics()
    #
    orientation = quadrotor_dynamics.run_simulation(pd_controller)

    # Plots
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(2, 3, figsize=(12, 6))
    axes[0, 0].plot(orientation[:, 0])
    axes[0, 0].set_title("$\omega_x$")
    axes[0, 1].plot(orientation[:, 1], label="$\omega_y$")
    axes[0, 1].set_title("$\omega_y$")
    axes[0, 2].plot(orientation[:, 2], label="$\omega_z$")
    axes[0, 2].set_title("$\omega_z$")
    axes[1, 0].plot(orientation[:, 3])
    axes[1, 0].set_title("$\dot{\omega}_x$")
    axes[1, 1].plot(orientation[:, 4], label="$\dot{\omega}_y$")
    axes[1, 1].set_title("$\dot{\omega}_y$")
    axes[1, 2].plot(orientation[:, 5], label="$\dot{\omega}_z$")
    axes[1, 2].set_title("$\dot{\omega}_z$")
    plt.tight_layout()
    plt.show()

if __name__ == '__main__':
    main()