-
Notifications
You must be signed in to change notification settings - Fork 0
/
prediction_sim.py
143 lines (113 loc) · 5.75 KB
/
prediction_sim.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
from data_loader import TestSet
import numpy as np
from torch.utils.data import DataLoader
import torch
from End2EndNet import End2EndNet
from scipy import integrate
from matplotlib import pyplot as plt
# prediction_sim.py: Simulate robotic system motion and predictive models over trajectory samples
def plot_flights(pred, label, image_count):
# Generate plots of the true and predicted flight paths
fig = plt.figure()
ax = plt.axes(projection='3d')
ax.plot3D(pred[:, 3], pred[:, 4], pred[:, 5], 'blue', linewidth=1)
ax.plot3D(label[:, 3], label[:, 4], label[:, 5], 'red', linewidth=1)
ax.legend(["Predicted Flight Path", "Actual Flight Path"])
ax.set_title('Flight Path')
ax.set_xlabel('X-position (m)')
ax.set_ylabel('Y-position (m)')
ax.set_zlabel('Z-position (m)')
plt.show()
fig.savefig('neural_net_sim{}.png'.format(image_count))
class NeuralNetModel:
# Quadrotor motion prediction model with Neural Net for quadrotor state prediction
def __init__(self, vels, init_state=np.zeros((6, 1))):
# Initialize variables to store quadrotor motion state
self.ang = init_state[0:3] # Angular position (XYZ Euler Angle)
self.pos = init_state[3:6] # Linear Position
self.step = 0
self.vels = vels # Neural network outputs full velocity and body rate tensor
# Initialize integration method for discrete time dynamics
self.ode = integrate.ode(self.state_dot).set_integrator('vode', nsteps=500, method='bdf')
def state_dot(self, time, state):
rate = self.vels[:3, self.step]
vel = self.vels[3:, self.step]
# Calculate rotation matrix from body frame to inertial frame
s_phi = np.sin(self.ang[0]).item()
c_phi = np.cos(self.ang[0]).item()
c_theta = np.cos(self.ang[1]).item()
# Calculate Euler angle time derivatives
M = np.asarray([[1, 0, -s_phi], [0, c_phi, s_phi * c_theta], [0, -s_phi, c_theta * c_phi]])
ang_dot = np.dot(np.linalg.inv(M), rate)
# Concatenate into final state derivative vector
state_dot = np.concatenate([ang_dot, vel])
return state_dot
def update(self, dt):
init_state = np.concatenate([self.ang, self.pos]) # Set initial state for integration
self.ode.set_initial_value(init_state, 0) # Initialize ODE
updated_state = self.ode.integrate(self.ode.t + dt) # Integrate from t to t+dt
self.ang = updated_state[0:3]
self.pos = updated_state[3:6]
if __name__ == "__main__":
# Define training/validation datasets and dataloaders
lookback = 64
pred_steps = 90
test_set = TestSet('data/AscTec_Pelican_Flight_Dataset.mat', lookback, pred_steps, full_state=True)
test_loader = DataLoader(test_set, batch_size=1, shuffle=True, num_workers=0)
if torch.cuda.is_available():
device = torch.device("cuda:0")
torch.set_default_tensor_type("torch.cuda.FloatTensor")
print("GPU")
else:
device = torch.device("cpu")
print("CPU")
# Load neural network model
vel_model = End2EndNet(lookback, pred_steps)
vel_model.load_state_dict(torch.load('./End2EndNet.pth', map_location=device))
vel_model.train(False)
vel_model.eval()
outliers = 0
count = 0
image_count = 0
outlier_range = np.zeros((16, 1))
inlier_range = np.zeros((16))
# Evaluate for entire test set
with torch.no_grad():
for data in test_loader:
raw_input = torch.transpose(data["input"].type(torch.FloatTensor), 1, 2).to(device) # Load Input data
label = torch.transpose(data["label"].type(torch.FloatTensor), 1, 2).to(device) # Load labels
feedforward = torch.zeros(label.shape) # Future control inputs
feedforward[:, 12:, :] = label[:, 12:, :]
input = torch.cat((raw_input, feedforward), 2) # Full model input
pred_vels = vel_model(input) # Neural network truncated state prediction
init_vel = torch.zeros((1, 6, 1))
init_vel[:, :, 0] = raw_input[:, 6:12, -1]
vels = torch.cat((init_vel, pred_vels), 2)
# Forward integration of predicted velocities
black_box_model = NeuralNetModel(vels[0, :, :].numpy(), raw_input[0, :6, -1].numpy())
state_rec = np.zeros((pred_steps+1, 12))
state_rec[:, 6:] = black_box_model.vels.T
for i in range(pred_steps):
state_rec[i, 0:3] = black_box_model.ang
state_rec[i, 3:6] = black_box_model.pos
black_box_model.update(0.01)
black_box_model.step += 1
state_rec[pred_steps, 0:3] = black_box_model.ang
state_rec[pred_steps, 3:6] = black_box_model.pos
# Calculate errors
state_label = label[0, :12, :].numpy().T
state_error = (state_rec[:pred_steps, :] - state_label) ** 2
full_sequence = np.concatenate((raw_input.numpy(), label.numpy()), axis=2)
# Calculate quadrotor state range for all samples
inlier_range += full_sequence[0, :, :].ptp(axis=1)/len(test_set)
# Calculate quadrotor state range for high error samples
if state_error.sum() > 50:
ranges = np.expand_dims(full_sequence[0, :, :].ptp(axis=1), axis=1)
outlier_range = np.concatenate((outlier_range, ranges), axis=1)
# Calculate quadrotor state range for low error samples
if state_error.sum() < 1:
image_count += 1
plot_flights(state_rec[:pred_steps, :], state_label, image_count)
count += 1
if count % 10 == 0:
print(count)