forked from mbrossar/SE2-3-
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathintro.py
145 lines (129 loc) · 5.47 KB
/
intro.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import torch
from utils import *
from lie_group_utils import SO3, SE3_2
from preintegration_utils import *
import matplotlib.pyplot as plt
plt.rcParams["legend.loc"] = "upper right"
plt.rcParams['axes.titlesize'] = 'x-large'
plt.rcParams['axes.labelsize'] = 'x-large'
plt.rcParams['legend.fontsize'] = 'x-large'
plt.rcParams['xtick.labelsize'] = 'x-large'
plt.rcParams['ytick.labelsize'] = 'x-large'
plt.rcParams['text.usetex'] =True
params= {'text.latex.preamble' : [r'\usepackage{amsmath}',
r'\usepackage{amssymb}']}
plt.rcParams.update(params)
import numpy as np
import scipy.linalg
torch.set_default_dtype(torch.float64)
def propagate(T0, Sigma, Upsilon, Q, method, dt, g, cholQ=0):
"""
Propagate an extended pose and its uncertainty
"""
Gamma = f_Gamma(g, dt)
Phi = f_flux(T0, dt)
# compound the mean
T = Gamma.mm(Phi).mm(Upsilon)
# Jacobian for propagating prior along time
F = torch.eye(9)
F[6:9, 3:6] = torch.eye(3)*dt
# compute Adjoint of right transformation mean
AdUps = SE3_2.uAd(SE3_2.uinv(Upsilon))
Sigma_tmp = axat(AdUps.mm(F), Sigma)
# compound the covariances based on the second-order method
Sigma_prop = Sigma_tmp + Q
if method == 2:
# SO(3) x R^6
wedge_acc = SO3.uwedge(Upsilon[:3, 3]) # already multiplied by dt
F = torch.eye(9)
F[3:6, :3] = T0[:3, :3].t()
F[3:6, :3] = -T0[:3, :3].mm(wedge_acc)
F[6:9, :3] = F[3:6, :3]*dt/2
F[6:9, 3:6] = dt*torch.eye(3)
G = torch.zeros(9, 6)
G[:3, :3] = dt*T0[:3, :3].t()
G[3:6, 3:6] = T0[:3, :3]*dt
G[6:9, 3:6] = 1/2*T0[:3, :3]*(dt**2)
Sigma_prop = axat(F, Sigma) + axat(G, Q[:6, :6]/(dt**2))
Sigma_prop = (Sigma_prop + Sigma_prop.t())/2
return T, Sigma_prop
def plot_se23_helper(T_est, P_est, color, i1, i2, i_max, path):
P_est_chol = torch.cholesky(P_est + torch.eye(9)*1e-16)
r = bmv(P_est_chol.expand(i_max, 9, 9), torch.randn(i_max, 9))
Ttemp = T_est[-1].expand(i_max, 5, 5).bmm(SE3_2.exp(r.cuda()).cpu())
p_est = Ttemp[:, :3, 4]
plt.scatter(p_est[:, i1], p_est[:, i2], color=color, s=3)
# if i1 == 0 and i2 == 1:
# np.savetxt(path, p_est.numpy(), header="x y z", comments='')
def plot_so3_helper(T_est, P_est, color, i1, i2, i_max, path):
r = torch.randn(i_max, 3)
P_est_chol = torch.cholesky(P_est[6:9, 6:9]+torch.eye(3)*1e-16)
p_est = bmv(P_est_chol.expand(i_max, 3, 3), r) + T_est[-1, :3, 4]
plt.scatter(p_est[:, i1], p_est[:, i2], color=color, s=2, alpha=0.5)
# if i1 == 0 and i2 == 1:
# np.savetxt(path, p_est.numpy(), header="x y z", comments='')
def main(i_max, k_max, T0, Sigma0, Upsilon, Q, cholQ, dt, g, sigma, m_max, paths):
# Generate some random samples
T = torch.zeros(i_max, k_max, 5, 5).cuda()
T[:, 0] = T0.cuda().repeat(i_max, 1, 1)
# NOTE: no initial uncertainty
Gamma = f_Gamma(g, dt).cuda().expand(i_max, 5, 5)
tmp = cholQ.cuda().expand(i_max, 9, 9)
tmp2 = Upsilon.cuda().expand(i_max, 5, 5)
for k in range(1, k_max):
T_k = SE3_2.exp(bmv(tmp, torch.randn(i_max, 9).cuda()))
Phi = f_flux(T[:, k-1], dt)
T[:, k] = Gamma.bmm(Phi).bmm(T_k).bmm(tmp2)
T = T.cpu()
# Propagate the uncertainty methods
T_est = torch.zeros(k_max, 5, 5)
P_est = torch.zeros(k_max, 9, 9)
SigmaSO3 = torch.zeros(k_max, 9, 9) # SO(3) x R^6 covariance
T_est[0] = T0
P_est[0] = Sigma0.clone()
SigmaSO3[0] = Sigma0.clone()
for k in range(1, k_max):
T_est[k], P_est[k] = propagate(T_est[k-1], P_est[k-1], Upsilon, Q, 1, dt, g)
# baseline method
_, SigmaSO3[k] = propagate(T_est[k-1], SigmaSO3[k-1], Upsilon, Q, 2, dt, g)
# Now plot the transformations
labels = ['x (m)', 'y (m)', 'z (m)']
for i1, i2 in ((0, 1), (0, 2), (1, 2)):
plt.figure()
# Plot the covariance of the samples
plot_so3_helper(T_est, SigmaSO3[-1], "red", i1, i2, i_max, paths[1])
# Plot the propagated covariance projected onto i1, i2
plot_se23_helper(T_est, P_est[-1], 'green', i1, i2, i_max, paths[2])
# Plot the random samples' xy-locations
plt.scatter(T[:, -1, i1, 4], T[:, -1, i2, 4], s=1, color='black', alpha=0.5)
plt.scatter(T_est[-1, i1, 4], T_est[-1, i2, 4], color='yellow', s=30)
plt.xlabel(labels[i1])
plt.ylabel(labels[i2])
plt.legend([r"$SO(3) \times \mathbb{R}^6$", r"$SE_2(3)$"])
# np.savetxt(paths[0], T[:, -1, :3, 4].numpy(), header="x y z", comments='')
plt.show()
if __name__ == '__main__':
### Parameters ###
i_max = 1200 # number of random points
k_max = 301 # number of compounded poses
sigma = 3 # plot option (factor for ellipse size)
m_max = 100 # plot option (number of points for each ellipse)
g = torch.Tensor([0, 0, 9.81]) # gravity vector
dt = 0.01 # step time (s)
paths = ["figures/intro_T.txt", "figures/intro_est.txt", "figures/intro_b.txt"]
# Define a PDF over transformations
Upsilon = torch.eye(5)
Upsilon[:3, 3] = -g*dt
Upsilon[:3, 4] = Upsilon[:3, 3]*dt/2
# Define right perturbation noise
cholQ = Upsilon.new_zeros(9, 9)
cholQ[:3, :3] = 1e-2*torch.eye(3)
cholQ[3:6, 3:6] = 5e-3*torch.eye(3)
cholQ[6:9, 3:6] = cholQ[3:6, 3:6]*dt/2
Q = cholQ.mm(cholQ.t())
# initial state
T0 = torch.eye(5)
T0[:3, 3] = torch.Tensor([5, 0, 0])
# initial right perturbation uncertainty
Sigma0 = torch.zeros(9, 9)
main(i_max, k_max, T0, Sigma0, Upsilon, Q, cholQ, dt, g, sigma, m_max, paths)