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]])
