MPC - cost on \delta u #627

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import control as ct
import control.optimal as opt

if __name__ == '__main__':

    # Set seed for reproducibility
    np.random.seed(42)

    # Settings
    n_sim = 100

    # References
    p_ref = 7.0
    v_ref = 0.0
    x_ref = np.array([p_ref, v_ref])  # reference state
    u_ref = 0.0
    du_min = -0.5
    du_max = 0.5

    # Constants #
    Ts = 0.2  # sampling time (s)
    M = 2     # mass (Kg)
    b = 0.3   # friction coefficient (N*s/m)

    # Initial input and state
    u_0 = np.array([0.0])
    x_0 = np.array([0.1, 0.2])

    # Continuous-time system matrices
    Ac = np.array([
        [0.0, 1.0],
        [0, -b/M]]
    )
    Bc = np.array([
        [0.0],
        [1/M]
    ])

    [n_x, n_u] = Bc.shape  # number of states and number or inputs

    # Cost function
    n_p = 25  # prediction horizon
    Q_x = np.diag([0.5, 0.05])  # cost on x
    Q_u = np.diag([0.0])  # cost on u
    Q_du = np.diag([0.1])  # cost on delta_u

    # Forward euler discretization
    Ad = np.eye(n_x) + Ac*Ts
    Bd = Bc*Ts

    # Augmented model as ct.NonlinearIOSystem, for compatibility with python-control
    def aug_model_update(t, x, delta_u, params={}):
        u = np.atleast_1d(x[0])
        x_sys = x[1:]
        x_new = np.copy(x)
        x_new[1:] = Ad@x_sys + Bd@u
        x_new[0] = x_new[0] + delta_u
        return x_new

    def aug_model_out(t, x, u, params={}):
        return x  # need full output for MPC...

    aug_model = ct.NonlinearIOSystem(aug_model_update, aug_model_out, name='aug_model', inputs='delta_force',
                                     outputs='position', states=('force', 'position', 'velocity'), dt=Ts)

    # True system as ct.NonlinearIOSystem, for compatibility with python-control
    def sys_update(t, x, u, params={}):
        x_new = Ad @ x + Bd@u
        return x_new

    def sys_out(t, x, u, params={}):
        return x  # need full output for MPC...

    sys_real = ct.NonlinearIOSystem(sys_update, sys_update, name='true_system', inputs='force',
                                    outputs='position', states=('position', 'velocity'), dt=Ts)

    # Set-up MPC problem
    constraints = [opt.input_range_constraint(aug_model, du_min, du_max)]
    cost = opt.quadratic_cost(aug_model, sp.linalg.block_diag(Q_u, Q_x), Q_du, x0=np.r_[u_ref, x_ref])
    ctrl = opt.create_mpc_iosystem(aug_model, np.arange(n_p) * Ts, cost, constraints, dt=aug_model.dt)

    # System input and state variables to be updated in the simulation loop
    u_sys = np.copy(u_0)
    x_sys = np.copy(x_0)
    x_mpc = np.zeros(n_p * n_u)  # MPC state: flattened mpc MV sequence

    # Simulation loop
    x_lst = []
    u_lst = []
    du_lst = []
    t_sim = np.arange(n_sim)*Ts
    for i in range(n_sim):
        # print(i)

        # Save current state and input
        x_lst.append(x_sys)
        u_lst.append(u_sys)

        # Controller step: compute optimal input (variation)
        x_mpc = ctrl.dynamics(None, x_mpc, np.r_[u_sys, x_sys])  # MPC system is fed by the system's augmented output
        du = ctrl.output(None, x_mpc, None)
        du_lst.append(du)

        # System (+ integrated input) update step
        x_sys = sys_real.dynamics(None, x_sys, u_sys)
        u_sys = u_sys + du

    # Stack everything in matrices
    x_mat = np.stack(x_lst)
    du_mat = np.stack(du_lst)
    u_mat = np.stack(u_lst)

    fig, axes = plt.subplots(3, 1, figsize=(10, 10))
    axes[0].plot(t_sim, x_mat[:, 0], "k", label='p')
    axes[0].plot(t_sim, x_ref[0]*np.ones_like(t_sim), "r--", label="p_ref")
    axes[0].set_title("Position (m)")

    axes[1].plot(t_sim, x_mat[:, 1], label="v")
    axes[1].plot(t_sim, x_ref[1]*np.ones_like(t_sim), "r--", label="v_ref")
    axes[1].set_title("Velocity (m/s)")

    axes[2].plot(t_sim, u_mat[:, 0], label="u")
    axes[2].plot(t_sim, u_ref*np.ones_like(t_sim), "r--", label="u_ref")
    axes[2].set_title("Force (N)")

    for ax in axes:
        ax.grid(True)
        ax.legend()