def mean(self, x):
f = self.gp
if self.grad_check:
old_dtype = self.dtype
self.to(torch.float64)
with variable_required_grad(x):
torch.autograd.gradcheck(f.mean, x.double())
self.to(dtype=old_dtype)
with variable_required_grad(x):
return torch.autograd.grad(f.mean(x), x)[0]
def covar(self, G, x, xp):
"""
returns covar(∇f, g) given covar(f, g)
"""
f = self.gp
with variable_required_grad(x):
J_covar_fg = t_jac(self.gp.covar(G, x, xp), x)
return J_covar_fg.t()
def knl(self, x, xp, eigeps=EPS):
f = self.gp
if xp is x:
xp = xp.detach().clone()
grad_k_func = lambda xs, xt: torch.autograd.grad(
f.knl(xs, xt), xs, create_graph=True)[0]
if self.grad_check:
old_dtype = self.dtype
self.to(torch.float64)
f_knl_func = lambda xt: f.knl(xt, xp.double())
with variable_required_grad(x):
torch.autograd.gradcheck(f_knl_func, x.double())
torch.autograd.gradgradcheck(lambda x: f.knl(x, x), x.double())
with variable_required_grad(x):
with variable_required_grad(xp):
torch.autograd.gradcheck(
lambda xp: grad_k_func(x.double(), xp)[0], xp.double())
self.to(dtype=old_dtype)
analytical = self.analytical_hessian
if analytical:
Hxx_k = t_hessian(f.knl, x, xp)
else:
with variable_required_grad(x):
with variable_required_grad(xp):
old_dtype = self.dtype
self.to(torch.float64)
Hxx_k = tgradcheck.get_numerical_jacobian(
partial(grad_k_func, x.double()), xp.double())
self.to(dtype=old_dtype)
Hxx_k = Hxx_k.to(old_dtype)
if torch.allclose(x, xp):
eigenvalues, eigenvectors = torch.eig(Hxx_k, eigenvectors=False)
assert (eigenvalues[:, 0] >
-eigeps).all(), " Hessian must be positive definite"
small_neg_eig = ((eigenvalues[:, 0] > -eigeps) &
(eigenvalues[:, 0] < 0))
if small_neg_eig.any():
eigenvalues, eigenvectors = torch.eig(Hxx_k, eigenvectors=True)
evalz = eigenvalues[:, 0]
evalz[small_neg_eig] = 0
Hxx_k = eigenvectors.T @ torch.diag(evalz) @ eigenvectors
return Hxx_k
def grad_h_col(self, X_in):
if X_in.ndim == 1:
X = X_in.unsqueeze(0)
with variable_required_grad(X):
grad_h_x = torch.autograd.grad(self.value(X), X)[0]
if X_in.ndim == 1:
grad_h_x = grad_h_x.squeeze(0)
return grad_h_x
def control(self, x, t=None):
from bdlqr.lqr import LinearSystem
with variable_required_grad(x) as xp:
A = t_jac(self.model.f_func(xp), xp)
B = to_numpy(self.model.g_func(x)) * self.dt
Q = self.Q
R = self.R
s=(- Q @ to_numpy(self.x_goal))
Q_T=Q
s_T=(- Q @ to_numpy(self.x_goal))
z=np.zeros(R.shape[-1])
T=(self.numSteps-t)
sys = LinearSystem(A=(np.eye(A.shape[-1]) + to_numpy(A) * self.dt),
B=B,
Q=Q,
R=R,
s=s,
Q_T=Q_T,
s_T=s_T,
z=z,
T=T
)
xs, us = sys.solve(to_numpy(x), 1)
with variable_required_grad(x) as xp:
A = t_jac(self.model.f_func(xp)
+ self.model.g_func(xp) @ torch.tensor(us[0], dtype=xp.dtype),
xp)
sys = LinearSystem(A=(np.eye(A.shape[-1]) + to_numpy(A) * self.dt),
B=B,
Q=Q,
R=R,
s=s,
Q_T=Q_T,
s_T=s_T,
z=z,
T=T
)
xs, us = sys.solve(to_numpy(x), 1)
return torch.tensor(us[0], dtype=xp.dtype)
def test_gradient_f_gp(dynamic_models, skip_test=False, dist=1e-4):
learned_model, true_model, xtest, utest = dynamic_models
grad_f = GradientGP(DeterministicGP(
lambda x: torch.tensor([1., 0.]), shape=(2, ), name="[1, 0]").t()
@ learned_model.fu_func_gp(utest),
x_shape=(2, ))
def xdot_func(x):
return true_model.f_func(x)[0] + (true_model.g_func(x) @ utest)[0]
with variable_required_grad(xtest):
true_grad_f = torch.autograd.grad(xdot_func(xtest), xtest)[0]
if not skip_test:
assert to_numpy(grad_f.mean(xtest)) == pytest.approx(
to_numpy(true_grad_f), abs=0.1, rel=0.4)
grad_f.knl(xtest, xtest)
return grad_f
def test_gradient_gp(dynamic_models,
skip_test=False,
dist=1e-4,
grad_check=True):
learned_model, true_model, xtest, _ = dynamic_models
if grad_check:
learned_model.double_()
func = lambda lm, X: lm.f_func_knl(X, xtest.double())[0, 0]
with variable_required_grad(xtest):
torch.autograd.gradcheck(partial(func, learned_model),
xtest.double())
learned_model.float_()
l1h = test_affine_gp(dynamic_models, skip_test=True)
true_cbf2 = RadialCBFRelDegree2(true_model)
grad_l1h = GradientGP(l1h, x_shape=xtest.shape)
if not skip_test:
assert to_numpy(grad_l1h.mean(xtest)) == pytest.approx(to_numpy(
true_cbf2.grad_lie_f_cbf(xtest)),
abs=0.1,
rel=0.4)
grad_l1h.knl(xtest, xtest)
return grad_l1h, l1h
def A(self, x, u, t=0):
with variable_required_grad(x) as xp:
dA = t_jac(self.f(xp, u, t=t), xp)
return torch.eye(x.shape[-1]) + dA * self.dt
def grad_lie_f_h_col(self, X):
with variable_required_grad(X):
return torch.autograd.grad(self.lie_f_h_col(X), X)[0]
def test_GP_train_predict(n=2,
m=3,
D=50,
deterministic=False,
rel_tol=0.10,
abs_tol=0.80,
perturb_scale=0.1,
sample_generator=sample_generator_trajectory,
dynamics_model_class=RandomDynamicsModel,
training_iter=100,
grad_predict=False):
if grad_predict:
deterministic = True
chosen_seed = torch.randint(100000, (1, ))
#chosen_seed = 52648
print("Random seed: {}".format(chosen_seed))
torch.manual_seed(chosen_seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Collect training data
dynamics_model = dynamics_model_class(m, n, deterministic=deterministic)
Xdot, X, U = sample_generator(dynamics_model, D)
if X.shape[-1] == 2 and U.shape[-1] == 1:
plot_results(torch.arange(U.shape[0]),
omega_vec=X[:-1, 0],
theta_vec=X[:-1, 1],
u_vec=U[:, 0])
# Test train split
shuffled_order = np.arange(D)
#shuffled_order = torch.randint(D, size=(D,))
np.random.shuffle(shuffled_order)
shuffled_order = torch.from_numpy(shuffled_order)
train_indices = shuffled_order[:int(D * 0.8)]
test_indices = shuffled_order[int(D * 0.8):]
# Train data
Xtrain, Utrain, XdotTrain = [Mat[train_indices, :] for Mat in (X, U, Xdot)]
UHtrain = torch.cat((Utrain.new_ones((Utrain.shape[0], 1)), Utrain), dim=1)
# Test data
Xtest, Utest, XdotTest = [Mat[test_indices, :] for Mat in (X, U, Xdot)]
# Call the training routine
dgp = ControlAffineRegressor(Xtrain.shape[-1], Utrain.shape[-1])
dgp_exact = ControlAffineRegressorExact(Xtrain.shape[-1], Utrain.shape[-1])
dgp_vector = ControlAffineRegressorVector(Xtrain.shape[-1],
Utrain.shape[-1])
# Test prior
_ = dgp.predict(Xtest, return_cov=False)
dgp._fit_with_warnings(Xtrain,
Utrain,
XdotTrain,
training_iter=training_iter,
lr=0.01)
_, _ = dgp_exact.custom_predict(Xtest, compute_cov=False)
dgp_exact._fit_with_warnings(Xtrain,
Utrain,
XdotTrain,
training_iter=training_iter,
lr=0.01)
_, _ = dgp_vector.custom_predict(Xtest, compute_cov=False)
dgp_vector._fit_with_warnings(Xtrain,
Utrain,
XdotTrain,
training_iter=training_iter,
lr=0.01)
if X.shape[-1] == 2 and U.shape[-1] == 1 and PLOTTING:
plot_learned_2D_func(Xtrain.detach().cpu().numpy(), dgp.f_func,
dynamics_model.f_func)
plt.savefig('/tmp/f_learned_vs_f_true.pdf')
plot_learned_2D_func(Xtrain.detach().cpu().numpy(),
dgp.g_func,
dynamics_model.g_func,
axtitle="g(x)[{i}]")
plt.savefig('/tmp/g_learned_vs_g_true.pdf')
# check predicting training values
#FXT_train_mean, FXT_train_cov = dgp.predict(Xtrain)
#XdotGot_train = XdotTrain.new_empty(XdotTrain.shape)
#for i in range(Xtrain.shape[0]):
# XdotGot_train[i, :] = FXT_train_mean[i, :, :].T @ UHtrain[i, :]
predict_flatten_deprecated = True
if not predict_flatten_deprecated:
XdotGot_train_mean, XdotGot_train_cov = dgp._predict_flatten(
Xtrain[:-1], Utrain[:-1])
assert XdotGot_train_mean.detach().cpu().numpy() == pytest.approx(
XdotTrain[:-1].detach().cpu().numpy(), rel=rel_tol,
abs=abs_tol), """
Train data check using original flatten predict """
UHtest = torch.cat((Utest.new_ones((Utest.shape[0], 1)), Utest), dim=1)
if deterministic:
FXTexpected = torch.empty((Xtest.shape[0], 1 + m, n))
for i in range(Xtest.shape[0]):
FXTexpected[i, ...] = torch.cat((dynamics_model.f_func(
Xtest[i, :])[None, :], dynamics_model.g_func(Xtest[i, :]).T),
dim=0)
assert torch.allclose(XdotTest[i, :],
FXTexpected[i, :, :].T @ UHtest[i, :])
# check predicting train values
XdotTrain_mean = dgp.fu_func_mean(Utrain[:-1], Xtrain[:-1])
XdotTrain_mean_exact = dgp_exact.fu_func_mean(Utrain[:-1], Xtrain[:-1])
XdotTrain_mean_vector = dgp_vector.fu_func_mean(Utrain[:-1], Xtrain[:-1])
assert XdotTrain_mean.detach().cpu().numpy() == pytest.approx(
XdotTrain[:-1].detach().cpu().numpy(), rel=rel_tol, abs=abs_tol), """
Train data check using custom flatten predict """
assert XdotTrain_mean_exact.detach().cpu().numpy() == pytest.approx(
XdotTrain[:-1].detach().cpu().numpy(), rel=rel_tol, abs=abs_tol), """
Train data check using ControlAffineRegressorExact.custom_predict """
assert XdotTrain_mean_vector.detach().cpu().numpy() == pytest.approx(
XdotTrain[:-1].detach().cpu().numpy(), rel=rel_tol, abs=abs_tol), """
Train data check using ControlAffineRegressorExact.custom_predict """
if grad_predict and n == 1:
x0 = Xtrain[9:10, :].detach().clone()
u0 = Utrain[9:10, :].detach().clone()
#est_grad_fx = dgp.grad_fu_func_mean(x0, u0)
true_fu_func = lambda X: dynamics_model.f_func(
X) + dynamics_model.g_func(X).bmm(u0.unsqueeze(-1)).squeeze(-1)
with variable_required_grad(x0):
true_grad_fx = torch.autograd.grad(true_fu_func(x0), x0)[0]
with variable_required_grad(x0):
est_grad_fx_2 = torch.autograd.grad(dgp.fu_func_mean(u0, x0),
x0)[0]
assert to_numpy(est_grad_fx_2) == pytest.approx(to_numpy(true_grad_fx),
rel=rel_tol,
abs=abs_tol)
#assert to_numpy(est_grad_fx) == pytest.approx(to_numpy(true_grad_fx), rel=rel_tol, abs=abs_tol)
# Check predicting perturbed train values
Xptrain = Xtrain[:-1] * (
1 + torch.rand(Xtrain.shape[0] - 1, 1) * perturb_scale)
Uptrain = Utrain[:-1] * (
1 + torch.rand(Xtrain.shape[0] - 1, 1) * perturb_scale)
Xdot_ptrain = dynamics_model.f_func(Xptrain) + dynamics_model.g_func(
Xptrain).bmm(Uptrain.unsqueeze(-1)).squeeze(-1)
if not predict_flatten_deprecated:
XdotGot_ptrain_mean, XdotGot_ptrain_cov = dgp._predict_flatten(
Xptrain, Uptrain)
assert XdotGot_ptrain_mean.detach().cpu().numpy() == pytest.approx(
Xdot_ptrain.detach().cpu().numpy(), rel=rel_tol, abs=abs_tol), """
Perturbed Train data check using original flatten predict """
XdotGot_ptrain_mean_custom = dgp.fu_func_mean(Uptrain, Xptrain)
assert XdotGot_ptrain_mean_custom.detach().cpu().numpy() == pytest.approx(
Xdot_ptrain.detach().cpu().numpy(), rel=rel_tol, abs=abs_tol), """
Perturbed Train data check using custom flatten predict """
# check predicting test values
# FXTmean, FXTcov = dgp.predict(Xtest)
# XdotGot = XdotTest.new_empty(XdotTest.shape)
# for i in range(Xtest.shape[0]):
# XdotGot[i, :] = FXTmean[i, :, :].T @ UHtest[i, :]
if not predict_flatten_deprecated:
XdotGot_mean, XdotGot_cov = dgp.predict_flatten(Xtest, Utest)
assert XdotGot_mean.detach().cpu().numpy() == pytest.approx(
XdotTest.detach().cpu().numpy(), rel=rel_tol, abs=abs_tol)
#abs=XdotGot_cov.flatten().max())
# check predicting test values
Xdot_mean = dgp.fu_func_mean(Utest, Xtest)
Xdot_mean_exact = dgp_exact.fu_func_mean(Utest, Xtest)
Xdot_mean_vector = dgp_vector.fu_func_mean(Utest, Xtest)
assert Xdot_mean.detach().cpu().numpy() == pytest.approx(
XdotTest.detach().cpu().numpy(), rel=rel_tol, abs=abs_tol)
assert Xdot_mean_exact.detach().cpu().numpy() == pytest.approx(
XdotTest.detach().cpu().numpy(), rel=rel_tol, abs=abs_tol)
assert Xdot_mean_vector.detach().cpu().numpy() == pytest.approx(
XdotTest.detach().cpu().numpy(), rel=rel_tol, abs=abs_tol)
return dgp, dynamics_model