lqg for an inverted pendulum

lqg.py

"""
Inverted Pendulum LQG control
author: Edward Morgan
"""

import casadi as cs
import math
import time
from numpy.linalg import inv, eig
import matplotlib.pyplot as plt
import numpy as np

# Model parameters
l_bar = 2.0  # length of bar
M = 1.0  # [kg]
m = 0.3  # [kg]
g = 9.8  # [m/s^2]

nx = 4  # number of state
nu = 1  # number of input
Q = np.diag([.0, 1.0, 1.0, 0.0])  # state cost matrix
R = np.diag([0.01])  # input cost matrix


# agressive demo
# Q = np.diag([10.0, 10.0, 10.0, 10.0])  # state cost matrix
# R = np.diag([0.01])  # input cost matrix


delta_t = 0.1  # time tick [s]
sim_time = 15.0  # simulation time [s]

show_animation = True

def main():
    x0 = np.array([
        [0.0],
        [0.0],
        [0.4],
        [0.0]
    ]) # initial state of the robot

    Pk0 = np.array([[10,0,0,0],[0,10,0,0],[0,0,10,0],[0,0,0,10]]) # initial Covariance matrix
    Qk = np.diag([0.0, 0.0, 0.0, 0.0]) # process noise
    Rk = np.diag([0.1, 0.1]) # measurement noise
    N = 500

    x = np.copy(x0)
    x_est = np.copy(x0)
    time = 0.0

    kf_function, C = kalman_Filter()

    xhat = np.zeros((4,N+1))
    actual = np.zeros((4,N+1))
    Phat = np.zeros((4,N+1))
    kalman_err = np.zeros((2,N+1))
    t = np.zeros((1,N+1))
    u_ = np.zeros((2,N))

    
    xhat[:,0] = np.copy(x0).reshape(4,)
    Phat[:,0] = np.diag(Pk0)
    actual[:,0] = np.copy(x0).reshape(4,)
    Pk = Pk0
    # while sim_time > time:
    for k in range(N):
        t[:,k] = time
        

        # calc control input
        u = lqr_control(x_est)

        # simulate inverted pendulum cart
        x = simulation(x_est, u)

        noise = 0.01

        # y measurement
        x_measured =  np.copy(x)
        x_measured[0] = x[0] + noise*np.random.randn()
        x_measured[2] = x[2] + noise*np.random.randn()
        yt = (C@x_measured).reshape(-1)

        # kalman filter
        x_est, Pk, err_  = kf_function(x, Pk, Qk,  Rk, yt)

        # recording data for plotting
        u_[:, k] = u
        Phat[:, k+1] = np.diag(Pk)
        xhat[:,k+1] =  np.array(x_est).reshape(-1)
        kalman_err[:,k] = np.array(err_).reshape(-1)
        actual[:,k] = np.array(x_measured).reshape(-1)

        time += delta_t
        if show_animation:
            plt.clf()
            px = float(x[0, 0])
            theta = float(x[2, 0])
            plot_cart(px, theta)
            plt.xlim([-5.0, 2.0])
            plt.pause(0.001)

    t[:,k+1] = time
    print("Finish")
    print(u_.shape)
    print(f"x={float(x[0, 0]):.2f} [m] , theta={math.degrees(x[2, 0]):.2f} [deg]")
    results_plots(t.T, xhat, actual, Phat, u_)
    if show_animation:
        plt.show()




def simulation(x, u):
    A, B, C = get_model_matrix()
    x = A @ x + B @ u

    return x


def kalman_Filter():

    A, B,C = get_model_matrix()
    #kalman filter
    Qk = cs.MX.sym('Qk', 4,4) #process noise
    pka = cs.MX.sym('pka', 4,4)  #aposterior state covariance
    Rk = cs.MX.sym('Rk', 2,2) #measurements noise
    yk = cs.MX.sym('yk', 2,1) #measurements
    x = cs.MX.sym('xk', 4,1) #process

    # xk_ = A@x + B@u
    pk_ = A@pka@A.T + Qk

    #compute gain matrix
    L_k = Rk + C@pk_@C.T
    L_k_inv = cs.inv(L_k)
    K = pk_@C.T@L_k_inv

    #compute correction
    err_k = yk - C@x

    #new state estimate
    xk_p = x + K@err_k

    #propagate the estimation error covariance 
    Pk_p = (cs.MX.eye(x.size1()) - K@C)@pk_@(cs.MX.eye(x.size1()) - K@C).T + K@Rk@K.T

    xkhat_func = cs.Function('xkhat', [x, pka, Qk, Rk, yk], [xk_p, Pk_p, err_k]) #
    return xkhat_func, C


def solve_DARE(A, B, Q, R, maxiter=150, eps=0.01):
    """
    Solve a discrete time_Algebraic Riccati equation (DARE)
    """
    P = Q

    for i in range(maxiter):
        Pn = A.T @ P @ A - A.T @ P @ B @ \
            inv(R + B.T @ P @ B) @ B.T @ P @ A + Q
        if (abs(Pn - P)).max() < eps:
            break
        P = Pn

    return Pn


def dlqr(A, B, Q, R):
    """
    Solve the discrete time lqr controller.
    x[k+1] = A x[k] + B u[k]
    cost = sum x[k].T*Q*x[k] + u[k].T*R*u[k]
    # ref Bertsekas, p.151
    """

    # first, try to solve the ricatti equation
    P = solve_DARE(A, B, Q, R)

    # compute the LQR gain
    K = inv(B.T @ P @ B + R) @ (B.T @ P @ A)

    eigVals, eigVecs = eig(A - B @ K)
    return K, P, eigVals


def lqr_control(x):
    A, B, C = get_model_matrix()
    start = time.time()
    K, _, _ = dlqr(A, B, Q, R)
    u = -K @ x
    elapsed_time = time.time() - start
    print(f"calc time:{elapsed_time:.6f} [sec]")
    return u


def get_numpy_array_from_matrix(x):
    """
    get build-in list from matrix
    """
    return np.array(x).flatten()


def get_model_matrix():
    A = np.array([
        [0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, m * g / M, 0.0],
        [0.0, 0.0, 0.0, 1.0],
        [0.0, 0.0, g * (M + m) / (l_bar * M), 0.0]
    ])
    A = np.eye(nx) + delta_t * A

    B = np.array([
        [0.0],
        [1.0 / M],
        [0.0],
        [1.0 / (l_bar * M)]
    ])
    B = delta_t * B

    C = np.array([[1, 0, 0, 0],
     [0, 0, 1 ,0]])

    return A, B, C


def flatten(a):
    return np.array(a).flatten()


def plot_cart(xt, theta):
    cart_w = 1.0
    cart_h = 0.5
    radius = 0.1

    cx = np.array([-cart_w / 2.0, cart_w / 2.0, cart_w /
                   2.0, -cart_w / 2.0, -cart_w / 2.0])
    cy = np.array([0.0, 0.0, cart_h, cart_h, 0.0])
    cy += radius * 2.0

    cx = cx + xt

    bx = np.array([0.0, l_bar * math.sin(-theta)])
    bx += xt
    by = np.array([cart_h, l_bar * math.cos(-theta) + cart_h])
    by += radius * 2.0

    angles = np.arange(0.0, math.pi * 2.0, math.radians(3.0))
    ox = np.array([radius * math.cos(a) for a in angles])
    oy = np.array([radius * math.sin(a) for a in angles])

    rwx = np.copy(ox) + cart_w / 4.0 + xt
    rwy = np.copy(oy) + radius
    lwx = np.copy(ox) - cart_w / 4.0 + xt
    lwy = np.copy(oy) + radius

    wx = np.copy(ox) + bx[-1]
    wy = np.copy(oy) + by[-1]

    plt.plot(flatten(cx), flatten(cy), "-b")
    plt.plot(flatten(bx), flatten(by), "-k")
    plt.plot(flatten(rwx), flatten(rwy), "-k")
    plt.plot(flatten(lwx), flatten(lwy), "-k")
    plt.plot(flatten(wx), flatten(wy), "-k")
    plt.title(f"x: {xt:.2f} , theta: {math.degrees(theta):.2f}")

    # for stopping simulation with the esc key.
    plt.gcf().canvas.mpl_connect(
        'key_release_event',
        lambda event: [exit(0) if event.key == 'escape' else None])

    plt.axis("equal")


def results_plots(t, xhat_lqr, xreal_lqr, Phat_lqr, actions):
    w=50
    t = t.reshape(-1)

    plt.figure()
    plt.plot(t[:w],actions[0,:][:w], label="input at mobile base")
    plt.legend()

    plt.figure()
    plt.plot(t[:w],xhat_lqr[0,:][:w], label="estimated x[m]")
    plt.plot(t[:w],xreal_lqr[0,:][:w], label="real x[m]")
    # Adding uncertainty interval around the estimated values
    sigma_x = np.sqrt(Phat_lqr[0, :])  # Standard deviation is the sqrt of variance
    lower = xhat_lqr[0,:] - sigma_x
    upper = xhat_lqr[0,:] + sigma_x
    plt.fill_between(t[:w], lower[:w], upper[:w], color='gray', alpha=0.3, label="Uncertainty Interval")
    # plt.ylim(-0.1, 0.5)
    plt.legend()

    plt.figure()
    plt.plot(t[:w],xhat_lqr[1,:][:w], label="estimated dx[m/s]")
    plt.plot(t[:w],xreal_lqr[1,:][:w], label="real dx[m/s]")
    # Adding uncertainty interval around the estimated values
    sigma_xd = np.sqrt(Phat_lqr[1, :])  # Standard deviation is the sqrt of variance
    lower = xhat_lqr[1,:] - sigma_xd
    upper = xhat_lqr[1,:] + sigma_xd
    plt.fill_between(t[:w], lower[:w], upper[:w], color='gray', alpha=0.3, label="Uncertainty Interval")
    # plt.ylim(-0.2, 0.1)
    plt.legend()

    plt.figure()
    plt.plot(t[:w],xhat_lqr[2,:][:w], label="estimated thet[rad]")
    plt.plot(t[:w],xreal_lqr[2,:][:w], label="real thet[rad]")
    # Adding uncertainty interval around the estimated values
    sigma_thet = np.sqrt(Phat_lqr[2, :]) 
    lower = xhat_lqr[2,:] - sigma_thet
    upper = xhat_lqr[2,:] + sigma_thet
    plt.fill_between(t[:w], lower[:w], upper[:w], color='gray', alpha=0.3, label="Uncertainty Interval")
    plt.legend()
    # plt.ylim(-1, 3)

    plt.figure()
    plt.plot(t[:w],xhat_lqr[3,:][:w], label="estimated dthet[rad/s]")
    plt.plot(t[:w],xreal_lqr[3,:][:w], label="real dthet[rad/s]")
    # Adding uncertainty interval around the estimated values
    sigma_thetd = np.sqrt(Phat_lqr[3, :])  # Assuming Phat[3, 3, :] contains the variance for thet
    lower = xhat_lqr[3,:] - sigma_thetd
    upper = xhat_lqr[3,:] + sigma_thetd
    plt.fill_between(t[:w], lower[:w], upper[:w], color='gray', alpha=0.3, label="Uncertainty Interval")
    plt.legend()
    # plt.ylim(-2, 2)
    plt.show()

if __name__ == '__main__':
    main()