def check_phase(vector, dim=1, site=None):
'''
check the phase of one site translation
'''
if dim == 1:
new_vector = np.zeros_like(vector)
len_v = new_vector.size
new_vector[:int(len_v / 2)] = vector[::2]
new_vector[int(len_v / 2):] = vector[1::2]
return new_vector.conjugate().dot(vector)
if dim == 2:
new_vector = np.copy(vector)
len_v = new_vector.size
for i in range(site):
temp_vec = np.zeros_like(new_vector)
temp_vec[:int(len_v / 2)] = new_vector[::2]
temp_vec[int(len_v / 2):] = new_vector[1::2]
new_vector = np.copy(temp_vec)
return new_vector.conjugate().dot(vector)
def jnp_safe_copy(x):
try:
res = jnp.copy(x)
except NotImplementedError:
warn("`jnp.copy` is not implemented yet. " "Using the object's `copy` method.")
if hasattr(x, "copy"):
res = jnp.array(x.copy())
else:
warn(f"Object has no `copy` method: {x}")
res = x
return res
def LQR_solve(self, x0, u_init):
horizon = len(u_init)
x_array = self.forward_propagation(x0, u_init)
u_array = np.copy(u_init)
k_array, K_array = self.back_propagation(x_array, u_array)
res_dict = {
'x_array_star': np.array(x_array),
'u_array_star': np.array(u_array),
'k_array_opt': np.array(k_array),
'K_array_opt': np.array(K_array)
}
return res_dict
def finite_difference_second_order_(self, func, x):
n_dim = len(x)
func_x = func(x)
hessian = np.zeros((n_dim, n_dim))
for i in range(n_dim):
for j in range(n_dim):
x_copy = np.copy(x)
x_copy[i] += self.finite_diff_eps
x_copy[j] += self.finite_diff_eps
fpp = func(x_copy)
x_copy = np.copy(x)
x_copy[i] += self.finite_diff_eps
fp_ = func(x_copy)
x_copy = np.copy(x)
x_copy[j] += self.finite_diff_eps
f_p = func(x_copy)
x_copy = np.copy(x)
x_copy[i] -= self.finite_diff_eps
fn_ = func(x_copy)
x_copy = np.copy(x)
x_copy[j] -= self.finite_diff_eps
f_n = func(x_copy)
x_copy = np.copy(x)
x_copy[i] -= self.finite_diff_eps
x_copy[j] -= self.finite_diff_eps
fnn = func(x_copy)
hessian[i, j] = fpp - fp_ - f_p - f_n - fn_ + fnn
hessian = (hessian + 2 * func_x) / (2 * self.finite_diff_eps**2)
return hessian
def ilqr_iterate(self, x0, u_init, n_itrs=50, tol=1e-6, verbose=True):
#initialize the regularization term
self.reg = 1
#derive the initial guess trajectory from the initial guess of u
x_array = self.forward_propagation(x0, u_init)
u_array = np.copy(u_init)
#initialize current trajectory cost
J_opt = self.evaluate_trajectory_cost(x_array, u_init)
J_hist = [J_opt]
#iterates...
converged = False
for i in range(n_itrs):
k_array, K_array = self.back_propagation(x_array, u_array)
norm_k = np.mean(np.linalg.norm(k_array, axis=1))
#apply the control to update the trajectory by trying different alpha
accept = False
for alpha in self.alpha_array:
x_array_new, u_array_new = self.apply_control(
x_array, u_array, k_array, K_array, alpha)
#evaluate the cost of this trial
J_new = self.evaluate_trajectory_cost(x_array_new, u_array_new)
if J_new < J_opt:
#see if it is converged
if np.abs((J_opt - J_new) / J_opt) < tol:
#replacement for the next iteration
J_opt = J_new
x_array = x_array_new
u_array = u_array_new
converged = True
break
else:
#replacement for the next iteration
J_opt = J_new
x_array = x_array_new
u_array = u_array_new
#successful step, decrease the regularization term
#momentum like adaptive regularization
self.reg = np.max(
[self.reg_min, self.reg / self.reg_factor])
accept = True
if verbose:
print(
'Iteration {0}:\tJ = {1};\tnorm_k = {2};\treg = {3}'
.format(i + 1, J_opt, norm_k,
np.log10(self.reg)))
break
else:
#don't accept this
accept = False
J_hist.append(J_opt)
#exit if converged...
if converged:
if verbose:
print('Converged at iteration {0}; J = {1}; reg = {2}'.
format(i + 1, J_opt, self.reg))
break
#see if all the trials are rejected
if not accept:
#need to increase regularization
#check if the regularization term is too large
if self.reg > self.reg_max:
if verbose:
print(
'Exceeds regularization limit at iteration {0}; terminate the iterations'
.format(i + 1))
break
self.reg = self.reg * self.reg_factor
if verbose:
print(
'Reject the control perturbation. Increase the regularization term.'
)
#prepare result dictionary
res_dict = {
'J_hist': np.array(J_hist),
'x_array_opt': np.array(x_array),
'u_array_opt': np.array(u_array),
'k_array_opt': np.array(k_array),
'K_array_opt': np.array(K_array)
}
return res_dict
def iLQR_iteration(self, x0, u_init, n_itrs=50, tol=1e-6, verbose=False):
x_array = self.forward_propagation(x0, u_init)
u_array = np.copy(u_init)
J_opt = self.evaluate_trajectory_cost(x_array, u_init)
J_hist = [J_opt]
converged = False
for i in range(n_itrs):
k_array, K_array = self.back_propagation(x_array, u_array)
norm_k = np.mean(np.linalg.norm(k_array, axis=1))
accept = False
for alpha in self.alpha_array:
x_array_new, u_array_new = self.apply_control(
x_array, u_array, k_array, K_array, alpha)
J_new = self.evaluate_trajectory_cost(x_array_new, u_array_new)
if J_new < J_opt:
if np.abs((J_opt - J_new) / J_opt) < tol:
J_opt = J_new
x_array = x_array_new
u_array = u_array_new
converged = True
break
else:
J_opt = J_new
x_array = x_array_new
u_array = u_array_new
self.reg = np.max(
[self.reg_min, self.reg / self.reg_factor])
accept = True
if verbose:
print(
f'Iteration {i+1,}:\tJ = {J_opt:.3f};\tnorm_k = {norm_k:.3f};\treg = {np.log10(self.reg):.3f}'
)
break
else:
accept = False
J_hist.append(J_opt)
if converged:
if verbose:
print('Converged at iteration {0}; J = {1}; reg = {2}'.
format(i + 1, J_opt, self.reg))
break
if not accept:
#need to increase regularization
#check if the regularization term is too large
if self.reg > self.reg_max:
if verbose:
print(
'Exceeds regularization limit at iteration {0}; terminate the iterations'
.format(i + 1))
break
self.reg = self.reg * self.reg_factor
if verbose:
print(
'Reject the control perturbation. Increase the regularization term.'
)
res_dict = {
'J_hist': np.array(J_hist),
'x_array_opt': np.array(x_array),
'u_array_opt': np.array(u_array),
'k_array_opt': np.array(k_array),
'K_array_opt': np.array(K_array)
}
return res_dict
def build_ilqr_tracking_solver(self, ref_pnts, weight_mats):
#figure out dimension
self.T_ = len(ref_pnts)
self.n_dims_ = len(ref_pnts[0])
self.ref_array = np.copy(ref_pnts)
self.weight_array = [mat for mat in weight_mats]
#clone weight mats if there are not enough weight mats
for i in range(self.T_ - len(self.weight_array)):
self.weight_array.append(self.weight_array[-1])
#build dynamics, second-order linear dynamical system
self.A_ = np.eye(self.n_dims_*2)
self.A_[0:self.n_dims_, self.n_dims_:] = np.eye(self.n_dims_) * self.dt_
self.B_ = np.zeros((self.n_dims_*2, self.n_dims_))
self.B_[self.n_dims_:, :] = np.eye(self.n_dims_) * self.dt_
self.plant_dyn_ = lambda x, u, t, aux: np.dot(self.A_, x) + np.dot(self.B_, u)
#build cost functions, quadratic ones
def tmp_cost_func(x, u, t, aux):
err = x[0:self.n_dims_] - self.ref_array[t]
#autograd does not allow A.dot(B)
cost = np.dot(np.dot(err, self.weight_array[t]), err) + np.sum(u**2) * self.R_
if t > self.T_-1:
#regularize velocity for the termination point
#autograd does not allow self increment
cost = cost + np.sum(x[self.n_dims_:]**2) * self.R_ * self.Q_vel_ratio_
return cost
self.cost_ = tmp_cost_func
self.ilqr_ = pylqr.PyLQR_iLQRSolver(T=self.T_-1, plant_dyn=self.plant_dyn_, cost=self.cost_, use_autograd=self.use_autograd)
if not self.use_autograd:
self.plant_dyn_dx_ = lambda x, u, t, aux: self.A_
self.plant_dyn_du_ = lambda x, u, t, aux: self.B_
def tmp_cost_func_dx(x, u, t, aux):
err = x[0:self.n_dims_] - self.ref_array[t]
grad = np.concatenate([2*err.dot(self.weight_array[t]), np.zeros(self.n_dims_)])
if t > self.T_-1:
grad[self.n_dims_:] = grad[self.n_dims_:] + 2 * self.R_ * self.Q_vel_ratio_ * x[self.n_dims_, :]
return grad
self.cost_dx_ = tmp_cost_func_dx
self.cost_du_ = lambda x, u, t, aux: 2 * self.R_ * u
def tmp_cost_func_dxx(x, u, t, aux):
hessian = np.zeros((2*self.n_dims_, 2*self.n_dims_))
hessian[0:self.n_dims_, 0:self.n_dims_] = 2 * self.weight_array[t]
if t > self.T_-1:
hessian[self.n_dims_:, self.n_dims_:] = 2 * np.eye(self.n_dims_) * self.R_ * self.Q_vel_ratio_
return hessian
self.cost_dxx_ = tmp_cost_func_dxx
self.cost_duu_ = lambda x, u, t, aux: 2 * self.R_ * np.eye(self.n_dims_)
self.cost_dux_ = lambda x, u, t, aux: np.zeros((self.n_dims_, 2*self.n_dims_))
#build an iLQR solver based on given functions...
self.ilqr_.plant_dyn_dx = self.plant_dyn_dx_
self.ilqr_.plant_dyn_du = self.plant_dyn_du_
self.ilqr_.cost_dx = self.cost_dx_
self.ilqr_.cost_du = self.cost_du_
self.ilqr_.cost_dxx = self.cost_dxx_
self.ilqr_.cost_duu = self.cost_duu_
self.ilqr_.cost_dux = self.cost_dux_
return
