def __init__(self, X, Y, U, phi=observables.monomials(2), psi=observables.monomials(2), regressor='ols', is_generator=False, p_inv=True): self.X = X self.Y = Y self.U = U self.phi = phi self.psi = psi # Get number of data points self.N = self.X.shape[1] # Construct Phi and Psi matrices self.Phi_X = self.phi(X) self.Phi_Y = self.phi(Y) self.Psi_U = self.psi(U) # TODO: Create Koopman Generator Tensor if is_generator is True # Get dimensions self.dim_phi = self.Phi_X.shape[0] self.dim_psi = self.Psi_U.shape[0] # Make sure data is full rank # checkMatrixRank(self.Phi_X, "Phi_X") # checkMatrixRank(self.Phi_Y, "Phi_Y") # checkMatrixRank(self.Psi_U, "Psi_U") # Make sure condition numbers are small # checkConditionNumber(self.Phi_X, "Phi_X") # checkConditionNumber(self.Phi_Y, "Phi_Y") # checkConditionNumber(self.Psi_U, "Psi_U") # Build matrix of kronecker products between u_i and x_i for all 0 <= i <= N self.kronMatrix = np.empty([self.dim_psi * self.dim_phi, self.N]) for i in range(self.N): self.kronMatrix[:, i] = np.kron(self.Psi_U[:, i], self.Phi_X[:, i]) # Solve for M and B if regressor == 'rrr': self.M = estimate_L.rrr(self.kronMatrix.T, self.Phi_Y.T).T self.B = estimate_L.rrr(self.Phi_X.T, self.X.T) if regressor == 'sindy': self.M = estimate_L.SINDy(self.kronMatrix.T, self.Phi_Y.T).T self.B = estimate_L.SINDy(self.Phi_X.T, self.X.T) else: self.M = estimate_L.ols(self.kronMatrix.T, self.Phi_Y.T, p_inv).T self.B = estimate_L.ols(self.Phi_X.T, self.X.T, p_inv) # reshape M into tensor K self.K = np.empty([self.dim_phi, self.dim_phi, self.dim_psi]) for i in range(self.dim_phi): self.K[i] = self.M[i].reshape([self.dim_phi, self.dim_psi], order='F')
Phi_Y = np.empty((d_phi, N)) for i, y in enumerate(Y.T): Phi_Y[:, i] = phi(y) Psi_U = np.empty((d_psi, N)) for i, u in enumerate(U.T): Psi_U[:, i] = psi(u) #%% Build kronMatrix kronMatrix = np.empty((d_psi * d_phi, N)) for i in range(N): kronMatrix[:, i] = np.kron(Psi_U[:, i], Phi_X[:, i]) #%% Estimate M M = estimate_L.ols(kronMatrix.T, Phi_Y.T).T #%% Reshape M into K tensor K = np.empty((d_phi, d_phi, d_psi)) for i in range(d_phi): K[i] = M[i].reshape((d_phi, d_psi), order='F') def K_u(K, u): return np.einsum('ijz,z->ij', K, psi(u)) #%% Training error def l2_norm(true_state, predicted_state): error = true_state - predicted_state squaredError = np.power(error, 2)
Psi_U[:, i] = psi(u) # def phi(x): # return np.array([1, x[0], x[1], x[0]*x[1], x[0]**2, x[1]**2, x[2], x[0]*x[2], x[1]*x[2], x[2]**2]) XU = np.append(X, U, axis=0) d_phi_xu = phi(XU[:, 0]).shape[0] # 6 + (3-1) + 2 = 10 Phi_XU = np.empty((d_phi_xu, N)) for i, xu in enumerate(XU.T): Phi_XU[:, i] = phi(xu) Phi_XU_prime = Phi_XU[:, 1:] Phi_XU = Phi_XU[:, :-1] #%% Concatenated state and action Koopman operator Koopman_operator = estimate_L.ols(Phi_XU.T, Phi_XU_prime.T).T fun_koopman_operator = estimate_L.ols(Phi_XU.T, Y[:, :-1].T).T #%% Separate Koopman operator per action Koopman_operators = [] # np.empty((action_grid.shape[0],d_phi,d_phi)) for action in split_datasets: if np.array_equal(split_datasets[action]['phi_x'], [1]): Koopman_operators.append(np.zeros((d_phi, d_phi))) continue Koopman_operators.append( estimate_L.ols(split_datasets[action]['phi_x'].T, split_datasets[action]['phi_x_prime'].T).T) Koopman_operators = np.array(Koopman_operators, dtype=object) #%% Build kronMatrix
Phi_X = phi(X) Phi_Y = phi(Y) Psi_U = psi(U) dim_phi = Phi_X[:, 0].shape[0] dim_psi = Psi_U[:, 0].shape[0] N = X.shape[1] print(Phi_X.shape) #%% Build kronMatrix kronMatrix = np.empty((dim_psi * dim_phi, N)) for i in range(N): kronMatrix[:, i] = np.kron(Psi_U[:, i], Phi_X[:, i]) #%% Estimate M M = estimate_L.ols(kronMatrix.T, Phi_Y.T).T B = estimate_L.ols(Phi_X.T, X.T) #%% Reshape M into K tensor K = np.empty((dim_phi, dim_phi, dim_psi)) for i in range(dim_phi): K[i] = M[i].reshape((dim_phi, dim_psi), order='F') def K_u(K, u): psi_u = psi(u.reshape(-1, 1))[:, 0] return np.einsum('ijz,z->ij', K, psi_u) #%% Define cost function def cost(x, u):
# Psi_U[:,i] = psi(u[0])[:,0] dim_phi = Phi_X.shape[0] dim_psi = Psi_U.shape[0] print("Phi_X shape:", Phi_X.shape) print("Psi_U shape:", Psi_U.shape) #%% Estimate Koopman operator for each action Koopman_operators = np.empty((env.action_space.n, dim_phi, dim_phi)) for action in range(env.action_space.n): if np.array_equal(phi(X_data[action]), [1]): Koopman_operators[action] = np.zeros([dim_phi, dim_phi]) continue Koopman_operators[action] = estimate_L.ols( phi(X_data[action]).T, phi(Y_data[action]).T).T Koopman_operators = np.array(Koopman_operators) #%% Estimate extended state Koopman operator extended_koopman_operator = estimate_L.ols(Phi_XU[:, :-1].T, Phi_XU[:, 1:].T).T extended_B = estimate_L.ols(Phi_XU.T, XU.T) #%% Build kronMatrix kronMatrix = np.empty((dim_psi * dim_phi, N)) for i in range(N): kronMatrix[:, i] = np.kron(Psi_U[:, i], Phi_X[:, i]) #%% Estimate M and B matrices num_ranks = 20 - 1
for u in split_datasets[action]['x']: split_datasets[action]['phi_x'].append(phi(int(u))[:, 0]) split_datasets[action]['phi_x_prime'] = [] for u in split_datasets[action]['x_prime']: split_datasets[action]['phi_x_prime'].append(phi(int(u))[:, 0]) split_datasets[action]['phi_x'] = np.array(split_datasets[action]['phi_x']) split_datasets[action]['phi_x_prime'] = np.array( split_datasets[action]['phi_x_prime']) Psi_U = np.empty((d_psi, N)) for i, u in enumerate(U.T): Psi_U[:, i] = psi(int(u[0]))[:, 0] #%% Koopman operator for extended state extended_state_koopman_operator = estimate_L.ols(Phi_XU[:, :-1].T, Phi_XU[:, 1:].T).T extended_state_to_state = estimate_L.ols(Phi_XU.T, Phi_X.T).T estimated_P = np.empty_like(P) for enumerated_extended_state in enumerated_extended_states: estimated_P[:,enumerated_extended_state[0], enumerated_extended_state[1]] = \ np.reshape((extended_state_to_state @ extended_state_koopman_operator @ phi_xu(enumerated_extended_state)), (d_phi,)) #%% Separate Koopman operator per action separate_koopman_operators = [] for action in range(d_u): separate_koopman_operators.append( estimate_L.ols(split_datasets[action]['phi_x'], split_datasets[action]['phi_x_prime']).T) separate_koopman_operators = np.array(separate_koopman_operators) multi_koopman_estimated_P = np.empty_like(P)
def getPsiPhiMatrix(Psi_U, Phi_X): psiPhiMatrix = np.empty( (num_lifted_action_features * num_lifted_state_features, num_lifted_state_observations)) for i in range(num_lifted_state_observations): kron = np.kron(Psi_U[:, i], Phi_X[:, i]) psiPhiMatrix[:, i] = kron return psiPhiMatrix #%% psiPhiMatrix = getPsiPhiMatrix(Psi_U, Phi_X) print("PsiPhiMatrix shape:", psiPhiMatrix.shape) M = estimate_L.ols(psiPhiMatrix.T, getPhiMatrix(Y_opt).T).T print("M shape:", M.shape) assert M.shape == (num_lifted_state_features, num_lifted_state_features * num_lifted_action_features) K = np.empty((num_lifted_state_features, num_lifted_state_features, num_lifted_action_features)) for i in range(M.shape[0]): K[i] = M[i].reshape( (num_lifted_state_features, num_lifted_action_features), order='F') print("K shape:", K.shape) def K_u(u): return np.einsum('ijz,z->ij', K, psi(u))
Phi_Y = phi(Y) Psi_U = psi(U) dim_phi = Phi_X.shape[0] dim_psi = Psi_U.shape[0] print("Phi_X shape:", Phi_X.shape) print("Psi_U shape:", Psi_U.shape) #%% Build kronMatrix kronMatrix = np.empty((dim_psi * dim_phi, N)) for i in range(N): kronMatrix[:,i] = np.kron(Psi_U[:,i], Phi_X[:,i]) #%% Estimate M and B matrices M = estimate_L.ols(kronMatrix.T, Phi_Y.T, pinv=False).T _B = estimate_L.ols(Phi_X.T, X.T, pinv=False) #%% Reshape M into K tensor K = np.empty((dim_phi, dim_phi, dim_psi)) for i in range(dim_phi): K[i] = M[i].reshape((dim_phi,dim_psi), order='F') def K_u(K, u): if len(u.shape) == 1: u = u.reshape(-1,1) # assume transposing row vector into column vector # u must be column vector return np.einsum('ijz,z->ij', K, psi(u)[:,0]) #%% L2 Norm def l2_norm(true_state, predicted_state):
[2, 0, 0, 4, 0, 4], [0, 0.5, 1, 0, 1, 0], [0, 1, 2, 0, 2, 0]]) YT = np.array( [ [0.5, 1, 2, 1, 2, 1], [1, 2, 4, 2, 4, 2] ] ) print(np.array_equal((np.linalg.pinv([email protected])).T, np.linalg.pinv(([email protected]).T))) #%% Test Bhat_0 = np.linalg.pinv([email protected])@[email protected] Bhat_1 = estimate_L.ols(XT.T, YT.T) np.array_equal(Bhat_0, Bhat_1) # Can we try multiple initial conditions? #%% Create large sample def f(x, u): if u == 0: return x/2 return x*2 def psi(u): if u == 0: return np.array([[1], [0]]) return np.array([[0], [1]])