complete_controller.py

import os
# import rospy
import pinocchio as pin
import numpy as np
from scipy.interpolate import CubicSpline
import random


class ControllerV2:
    def __init__(self, dt) -> None:
        # Define state vector dimention and measurement dimension
        self.n = 7          # State dimension
        self.m = 5          # Measurement dimension

        # Robot specs
        self.b = 0.2022     # Wheel distance
        self.r = 0.0985     # Wheel radius

        # Define target linear and angular reference velocities
        self.v_r = 3e-1
        self.omega_r = 1e-1
        # Initialize reference trajectory
        self.x_r = 0
        self.y_r = 0
        self.theta_r = 0
        # Initialize torques
        self.tau = np.array([0, 0])  
        # Define time step
        self.dt = dt

        # Compute absolute URDF file path using the current directory
        self.urdfFile = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tiago.urdf')
        # Initialize Pinocchio robot model
        self.model = pin.buildModelFromUrdf(self.urdfFile, pin.JointModelFreeFlyer())
        # Initialize Pinocchio data
        self.data = self.model.createData()
        # Initialize robot configuration and velocity
        self.q = pin.neutral(self.model)
        self.v = pin.utils.zero(self.model.nv) 

        # Call Centroidal Composite Rigid Body Algorithm
        pin.ccrba(self.model, self.data, self.q, self.v)
        # Extract centroidal mass and centroidal inertia matrix
        m = sum([inertia.mass for inertia in self.model.inertias])
        I_zz = self.data.Ig.inertia[2, 2] 
        self.M = np.array([[m, 0, 0], [0, m, 0], [0, 0, I_zz]])    

        # Define controller parameters
        self.k1 = 1
        self.k2 = 15
        self.k3 = 1
        self.k4 = 10
        self.k5 = 10

        # Initialize state vector
        self.x_est = np.array([1, 1, np.pi/4, 0, 0, 0, 0]) + 0.01 * np.array([1, 1, 1, 1, 1, 1, 1]) * np.random.normal(0, 1, 1)
        # Initialize state covariance
        self.P = np.eye(self.n)

        # UKF parameters
        self.alpha = 1
        self.beta = 2
        self.kappa = 0
        self.lambda_ = self.alpha**2 * self.n # (self.n + self.kappa) - self.n

        # Weights for sigma points
        self.Wm = np.full(2 * self.n + 1, 1 / (2 * (self.n + self.lambda_)))
        self.Wc = np.full(2 * self.n + 1, 1 / (2 * (self.n + self.lambda_)))
        self.Wm[0] = self.lambda_ / (self.n + self.lambda_)
        self.Wc[0] = self.Wm[0] + (1 - self.alpha**2 + self.beta)

        # Process and measurement noise covariances
        self.Q_scale = 1e-5
        self.R_scale = 1e-5
        self.Q = np.eye(self.n) * self.Q_scale
        self.R = np.eye(self.m) * self.R_scale
    

    ##### CONTROLLER

    def command(self, x_robot) -> None:
        # Start with the predict phase to estimate the current state
        self.predict()
        # Update the state with the measurement
        self.update(x_robot)

        # Extract estimated variables
        x_e = self.x_est[0]
        y_e = self.x_est[1]
        theta_e = self.x_est[2]
        omega_Le = self.x_est[3]
        omega_Re = self.x_est[4]
        i_Le = self.x_est[5]
        i_Re = self.x_est[6]

        # REFERENCE TRAJECTORY: eq. 10
        # Use Euler method to integrate the trajectory derived from reference velocities
        # We update theta_r last for consistency with Euler method
        self.x_r += self.dt * self.v_r * np.cos(self.theta_r)
        self.y_r += self.dt * self.v_r * np.sin(self.theta_r)
        self.theta_r += self.dt * self.omega_r
        
        # TRACKING ERROR: eq. 11
        # Compute rotation matrix
        R_e = np.array([
            [np.cos(theta_e), np.sin(theta_e), 0],
            [-np.sin(theta_e), np.cos(theta_e), 0],
            [0, 0, 1]])
        # Normilize theta error
        theta_err = self.theta_r - theta_e
        norm_theta_err = np.arctan2(np.sin(theta_err), np.cos(theta_err))
        # Compute trajectory error
        e_t = np.array([self.x_r - x_e, self.y_r - y_e, norm_theta_err])
        # Compute tracking error
        e = R_e @ e_t.T
        
        # AUXILIAR VELOCITY: eq. 12
        omega = self.omega_r + self.v_r / 2 * (self.k3 * (e[1] + self.k3 * e[2]) + (np.sin(e[2]) / self.k2))
        v = self.v_r * np.cos(e[2]) - self.k3 * e[2] * omega + self.k1 * e[0]  

        # DESIRED VELOCITY: eq. 13
        # Compute T matrix and its inverse
        T = self.r / (2 * self.b) * np.array([
            [self.b * (1 - i_Le), self.b * (1 - i_Re)],
            [-2 * (1 - i_Le), 2 * (1 - i_Re)]])
        T_inv = np.linalg.inv(T)
        # Compute desired velocities ξd = [omega_Ld, omega_Rd] = T @ [v, omega].T
        omega_Ld, omega_Rd = T_inv @ np.array([v, omega])

        # BACKSTEPPING: eq. 7
        # TODO: compute derivative of omega_Ld, omega_Rd (ξd_dot)
        omega_Ld_dot = 0
        omega_Rd_dot = 0
        # Compute auxiliary input
        # Compute auxiliary input
        u = np.array([omega_Ld_dot, omega_Rd_dot]) - np.array([self.k4 * (omega_Le - omega_Ld), self.k5 * (omega_Re - omega_Rd)]) 

        # CONTROL INPUT: eq 4
        # Compute S matrix
        S = (1 / (2 * self.b)) * np.array([
            [self.b * self.r * (1 - i_Le) * np.cos(theta_e), self.b * self.r * (1 - i_Re) * np.cos(theta_e)],
            [self.b * self.r * (1 - i_Le) * np.sin(theta_e), self.b * self.r * (1 - i_Re) * np.sin(theta_e)],
            [-2 * self.r * (1 - i_Le), 2 * self.r * (1 - i_Re)]])
        # Compute B matrix
        B = np.array([
            [np.cos(theta_e), np.cos(theta_e)],
            [np.sin(theta_e), np.sin(theta_e)],
            [-self.b / 2, self.b / 2]])
        # Compute B_bar and M_bar
        B_bar = S.T @ B
        M_bar = S.T @ self.M @ S
        # Compute tau
        self.tau = np.linalg.inv(B_bar) @ M_bar @ u.T


    ###### UKF

    def f(self, x: np.ndarray) -> np.ndarray:
        # Extract estimated variables
        theta_e = x[2]
        omega_Le = x[3]
        omega_Re = x[4]
        i_Le = x[5]
        i_Re = x[6]
        
        # Compute S matrix
        S = (1 / (2 * self.b)) * np.array([
            [self.b * self.r * (1 - i_Le) * np.cos(theta_e), self.b * self.r * (1 - i_Re) * np.cos(theta_e)],
            [self.b * self.r * (1 - i_Le) * np.sin(theta_e), self.b * self.r * (1 - i_Re) * np.sin(theta_e)],
            [-2 * self.r * (1 - i_Le), 2 * self.r * (1 - i_Re)]])
        # Compute B matrix
        B = np.array([
            [np.cos(theta_e), np.cos(theta_e)],
            [np.sin(theta_e), np.sin(theta_e)],
            [-self.b / 2, self.b / 2]])
        # Compute B_bar and M_bar
        B_bar = S.T @ B
        M_bar = S.T @ self.M @ S
        # Compute ξd_dot
        omega_Le_dot, omega_Re_dot = np.linalg.inv(M_bar) @ B_bar @ self.tau.T
        
        # Compute T matrix
        T = self.r / (2 * self.b) * np.array([
            [self.b * (1 - i_Le), self.b * (1 - i_Re)],
            [-2 * (1 - i_Le), 2 * (1 - i_Re)]])
        # Compute v and omega (eq. 8)
        v, omega = T @ np.array([omega_Le, omega_Re])

        # Euler method to integrate
        x[0] += self.dt * v * np.cos(theta_e)
        x[1] += self.dt * v * np.sin(theta_e)
        x[2] += self.dt * omega
        x[3] += self.dt * omega_Le_dot
        x[4] += self.dt * omega_Re_dot
        x[5] += np.random.normal(0, np.sqrt(self.Q[5][5]))
        x[5] = np.clip(x[5], -0.5, 0.5)
        x[6] += np.random.normal(0, np.sqrt(self.Q[6][6]))
        x[6] = np.clip(x[6], -0.5, 0.5)
        
        return x

    def compute_sigma_points(self) -> np.array:
        # Calculate square root of P matrix using Cholesky Decomposition
        sqrt_P = np.linalg.cholesky((self.lambda_ + self.n) * self.P)

        # Initialize sigma points
        sigma_points = np.zeros((2 * self.n + 1, self.n))
        sigma_points[0, :] = self.x_est
        
        # Calculate sigma points
        for k in range(self.n):
            spread_vector = sqrt_P[:, k]
            sigma_points[k + 1, :] = self.x_est + spread_vector
            sigma_points[self.n + k + 1, :] = self.x_est - spread_vector
        
        return sigma_points

    def predict(self) -> None:
        # Predict the state using the UKF equations
        sigma_points = self.compute_sigma_points()
        sigma_points = np.apply_along_axis(self.f, 1, sigma_points)
        self.sigma_points = sigma_points

        # Compute the predicted state mean
        self.x_pred = np.sum(self.Wm[:, np.newaxis] * sigma_points, axis=0)
        
        # Compute the predicted state covariance
        self.P_pred = self.Q.copy()
        for i in range(2 * self.n + 1):
            y = sigma_points[i] - self.x_pred
            self.P_pred += self.Wc[i] * np.outer(y, y)

    def update(self, z_robot): 
        # Compute the predicted measurement mean
        Z = self.sigma_points[:, :self.m]
        # Compute z_pred
        z_pred = np.sum(self.Wm[:, np.newaxis] * Z, axis=0)
        
        # Compute the predicted measurement covariance
        P_zz = self.R.copy()
        for i in range(2 * self.n + 1):
            y = Z[i] - z_pred
            P_zz += self.Wc[i] * np.outer(y, y)
        
        # Compute the cross-covariance matrix
        P_xz = np.zeros((self.n, self.m))
        for i in range(2 * self.n + 1):
            P_xz += self.Wc[i] * np.outer(self.sigma_points[i] - self.x_pred, Z[i] - z_pred)
        
        # Compute the Kalman gain
        K = P_xz @ np.linalg.inv(P_zz)
        
        # Update the state with the measurement
        self.x_est = self.x_pred + K @ (z_robot - z_pred)
        self.P = self.P_pred - K @ P_zz @ K.T