python/LQR_probabilistic.py

+import numpy as np
+from math import factorial
+import matplotlib.pyplot as plt
+from scipy.linalg import solve_continuous_are
+from scipy.stats import multivariate_normal
+
+# Parameters
+# ===============================
+param = lambda: None # Lazy way to define an empty class in python
+param.dt = 1e-2  # Time step length
+param.nbDeriv = 2 # Order of the dynamical system
+param.nbVarPos = 2 # Number of position variable
+param.nbVar = param.nbVarPos * param.nbDeriv # Dimension of state vector
+param.nbData = 100  # Number of datapoints
+param.rfactor = 1e-7  # Control weight term
+
+R = np.eye((param.nbData-1) * param.nbVarPos) * param.rfactor  # Control cost matrix
+Q = np.zeros((param.nbVar * param.nbData, param.nbVar * param.nbData)) # Task precision for augmented state
+
+xd = np.zeros([param.nbVar, param.nbData])
+target = np.random.uniform(size=param.nbVarPos)
+xd[:, param.nbData-1] = np.concatenate((target, np.zeros(param.nbVarPos)))
+xd = xd.T.flatten()
+
+Q[param.nbVar*(param.nbData-1):param.nbVar*param.nbData,param.nbVar*(param.nbData-1):param.nbVar*param.nbData] = 10.0 * np.eye(param.nbVar)
+
+A1d = np.zeros((param.nbDeriv,param.nbDeriv))
+B1d = np.zeros((param.nbDeriv,1))
+
+for i in range(param.nbDeriv):
+    A1d += np.diag( np.ones(param.nbDeriv-i), i ) * param.dt**i * 1/factorial(i)
+    B1d[param.nbDeriv-i-1] = param.dt**(i+1) * 1/factorial(i+1)
+
+A = np.eye(param.nbVar)
+A[:param.nbVar, :param.nbVar] = np.kron(A1d, np.identity(param.nbVarPos))
+B = np.zeros((param.nbVar, param.nbVarPos))
+B[:param.nbVar] = np.kron(B1d, np.identity(param.nbVarPos))
+
+Su = np.zeros((param.nbVar*param.nbData, param.nbVarPos * (param.nbData-1))) 
+Sx = np.kron(np.ones((param.nbData,1)), np.eye(param.nbVar,param.nbVar))
+
+M = B
+for i in range(1,param.nbData):
+    Sx[i*param.nbVar:param.nbData*param.nbVar,:] = np.dot(Sx[i*param.nbVar:param.nbData*param.nbVar,:], A)
+    Su[param.nbVar*i:param.nbVar*i+M.shape[0],0:M.shape[1]] = M
+    M = np.hstack((np.dot(A,M),B)) # [0,nb_state_var-1]
+
+
+x0 = np.random.uniform(size=param.nbVar)
+mean = Sx @ x0 + Su @ np.linalg.inv(Su.T @ Q @ Su + R) @ Su.T @ Q @ (xd - Sx @ x0)
+cov = Su @ (Su.T @ Q @ Su + R) @ Su.T
+
+
+trajectories = multivariate_normal.rvs(mean=mean, cov=cov, size=100)
+
+plt.figure()
+for k in range(100):
+    trajectory = trajectories[k, :]
+    trajectory = trajectory.reshape((param.nbData, param.nbVar))
+    plt.plot(trajectory[:, 0], trajectory[:, 1], alpha=0.5)
+
+plt.scatter(x0[0], x0[1], color='green', label='Initial State')
+plt.scatter(target[0], target[1], color='red', label='Target State')
+plt.legend()
+plt.grid()
+plt.show()