In [1]:
import numpy as np
import scipy.linalg
import matplotlib.pyplot as plt
In [2]:
# state transition matrices for linear system: 
#     x(t+1) = A x (t) + B u(t)
A = np.array([[1,1],[0,1]])
B = np.array([[0],[1]])

# LQR quadratic cost per state
Q = np.array([[1,0],[0,0]])

# initial condition for system
z0 = -1 # initial position
v0 = 0  # initial velocity
x0 = np.vstack((z0,v0))

# number of time steps to simulate
T = 30
In [3]:
def lqr(A,B,Q,R):
    # Solves for the optimal infinite-horizon LQR gain matrix given linear system (A,B) 
    # and cost function parameterized by (Q,R)
    
    # solve DARE:
    M=scipy.linalg.solve_discrete_are(A,B,Q,R)

    # K=(B'MB + R)^(-1)*(B'MA)
    return np.dot(scipy.linalg.inv(np.dot(np.dot(B.T,M),B)+R),(np.dot(np.dot(B.T,M),A)))
In [4]:
def simulate_static_gain(A,B,K,x0,T):
    # simulates the linear system (A,B) with static control law
    # u(t) = K x(t)
    # from initial condition x0 for T time steps
    #
    # returns matrices u and x of control and state trajectories, respectively.
    # dim 1 of these matrices is indexed by time.
    x = x0
    u = [[0]]
    for t in range(T-1):
        u_t = -K.dot(x[:,-1])
        x_tp1 = A.dot(x[:,-1])+B.dot(u_t)
        x = np.hstack((x, x_tp1[:,None]))
        u = np.hstack((u,u_t[:,None]))
    return x,u
In [5]:
K_lo = lqr(A,B,Q,1)
(x_lo,u_lo) = simulate_static_gain(A,B,K_lo,x0,T)

K_hi = lqr(A,B,Q,100)
(x_hi,u_hi) = simulate_static_gain(A,B,K_hi,x0,T)

def plot_trajectories(x1,x2,ylabel):
    colors = [ '#2D328F', '#F15C19' ]
    label_fontsize = 18
    tick_fontsize = 14
    linewidth = 3

    T = x1.shape[0]
    t = np.arange(0,T)
    plt.plot(t,x1,color=colors[0],linewidth=linewidth,label='R=1')
    plt.plot(t,x2,color=colors[1],linewidth=linewidth,label='R=10')
        
    plt.xlabel('time',fontsize=label_fontsize)
    plt.ylabel(ylabel,fontsize=label_fontsize)
    plt.legend(fontsize=label_fontsize)
    plt.xticks(fontsize=tick_fontsize)
    plt.yticks(fontsize=tick_fontsize)

    plt.grid(True)
    plt.show()
    
plot_trajectories(x_lo[0,:],x_hi[0,:],'Position')
plot_trajectories(x_lo[1,:],x_hi[1,:],'Velocity')
plot_trajectories(u_lo.T,u_hi.T,'Control Action')