Skip to content
Snippets Groups Projects
mpc.py 5.18 KiB
Newer Older
  • Learn to ignore specific revisions
  • import numpy as np
    
    from guided_mrmp.utils import Roomba
    
    from guided_mrmp.optimizer import Optimizer
    
    
    np.seterr(divide="ignore", invalid="ignore")
    
    import cvxpy as opt
    
    
    class MPC:
        def __init__(self, model, T, DT, state_cost, final_state_cost, input_cost, input_rate_cost):
            """
    
            Args:
                vehicle ():
                T ():
                DT ():
                state_cost ():
                final_state_cost ():
                input_cost ():
                input_rate_cost ():
            """
            self.nx = 3  # number of state vars 
            self.nu = 2  # number of input/control vars
    
            if len(state_cost) != self.nx:
                raise ValueError(f"State Error cost matrix shuld be of size {self.nx}")
            if len(final_state_cost) != self.nx:
                raise ValueError(f"End State Error cost matrix shuld be of size {self.nx}")
            if len(input_cost) != self.nu:
                raise ValueError(f"Control Effort cost matrix shuld be of size {self.nu}")
            if len(input_rate_cost) != self.nu:
                raise ValueError(
                    f"Control Effort Difference cost matrix shuld be of size {self.nu}"
                )
    
            self.robot_model = model
            self.dt = DT
    
            # how far we can look into the future divided by our dt 
            # is the number of control intervals
            self.control_horizon = int(T / DT) 
    
            # Weight for the error in state
            self.Q = np.diag(state_cost)
    
            # Weight for the error in final state
            self.Qf = np.diag(final_state_cost)
    
            # weight for error in control
            self.R = np.diag(input_cost)
            self.P = np.diag(input_rate_cost)
    
    
            # Instantiate the optimizer
            self.optimizer = Optimizer(self.nx, self.nu, self.control_horizon, self.Q, self.Qf, self.R, self.P)
    
    
    rachelmoan's avatar
    rachelmoan committed
     
    
        def get_linear_model_matrices_roomba(self,x_bar,u_bar):
            """
            Computes the approximated LTI state space model x' = Ax + Bu + C
    
            Args:
                x_bar (array-like): State vector [x, y, theta]
                u_bar (array-like): Input vector [v, omega]
    
            Returns:
                A_lin, B_lin, C_lin: Linearized state-space matrices
            """
    
            x = x_bar[0]
            y = x_bar[1]
            theta = x_bar[2]
    
            v = u_bar[0]
            omega = u_bar[1]
    
            ct = np.cos(theta)
            st = np.sin(theta)
    
            # Initialize matrix A with zeros and fill in appropriate elements
            A = np.zeros((self.nx, self.nx))
            A[0, 2] = -v * st
            A[1, 2] = v * ct
    
            # Discrete-time state matrix A_lin
            A_lin = np.eye(self.nx) + self.dt * A
    
            # Initialize matrix B with zeros and fill in appropriate elements
            B = np.zeros((self.nx, self.nu))
            B[0, 0] = ct
            B[1, 0] = st
            B[2, 1] = 1
    
            # Discrete-time input matrix B_lin
            B_lin = self.dt * B
    
            # Compute the non-linear state update equation f(x, u)
            f_xu = np.array([v * ct, v * st, omega]).reshape(self.nx, 1)
    
            # Compute the constant vector C_lin
            C_lin = (self.dt * (f_xu - np.dot(A, x_bar.reshape(self.nx, 1)) - np.dot(B, u_bar.reshape(self.nu, 1))).flatten())
    
            return A_lin, B_lin, C_lin
        
        def get_linear_model_matrices(self, x_bar, u_bar):
            """
            Computes the approximated LTI state space model x' = Ax + Bu + C
    
            Args:
                x_bar (array-like):
                u_bar (array-like):
    
            Returns:
    
            """
    
            x = x_bar[0]
            y = x_bar[1]
            v = x_bar[2]
            theta = x_bar[3]
    
            a = u_bar[0]
            delta = u_bar[1]
    
            ct = np.cos(theta)
            st = np.sin(theta)
            cd = np.cos(delta)
            td = np.tan(delta)
    
            A = np.zeros((self.nx, self.nx))
            A[0, 2] = ct
            A[0, 3] = -v * st
            A[1, 2] = st
            A[1, 3] = v * ct
            A[3, 2] = v * td / self.robot_model.wheelbase
            A_lin = np.eye(self.nx) + self.dt * A
    
            B = np.zeros((self.nx, self.nu))
            B[2, 0] = 1
            B[3, 1] = v / (self.robot_model.wheelbase * cd**2)
            B_lin = self.dt * B
    
            f_xu = np.array([v * ct, v * st, a, v * td / self.robot_model.wheelbase]).reshape(
                self.nx, 1
            )
            C_lin = (
                self.dt
                * (
                    f_xu
                    - np.dot(A, x_bar.reshape(self.nx, 1))
                    - np.dot(B, u_bar.reshape(self.nu, 1))
                ).flatten()
            )
            return A_lin, B_lin, C_lin
    
        def step(self, initial_state, target, prev_cmd):
    
            A, B, C = self.get_linear_model_matrices_roomba(initial_state, prev_cmd)  # Use Roomba model
    
            # Use the Optimizer class to solve the optimization problem
            x_opt, u_opt = self.optimizer.solve(initial_state, target, prev_cmd, A, B, C, self.robot_model, self.dt)
    
            return x_opt, u_opt
    
        
    if __name__ == "__main__":
        # Example usage:
        dt = 0.1
    
        roomba = Roomba()
    
        Q = [20, 20, 20]  # state error cost
        Qf = [30, 30, 30]  # state final error cost
        R = [10, 10]  # input cost
        P = [10, 10]  # input rate of change cost
        mpc = MPC(roomba, 5, dt, Q, Qf, R, P)
        x_bar = np.array([0.0, 0.0, 0.0])
        u_bar = np.array([1.0, 0.1])
        A_lin, B_lin, C_lin = mpc.get_linear_model_matrices_roomba(x_bar, u_bar)
        print(A_lin)
        print(B_lin)
        print(C_lin)