-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_model.py
116 lines (89 loc) · 4.48 KB
/
test_model.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
import torch
import torch.nn.functional as F
from tqdm import tqdm
from data import PendulumDataset
import numpy as np
import scipy.linalg as la
import matplotlib.pyplot as plt
from networks.lstm_autoencoder import LSTMAutoEncoder
from networks.linear_autoencoder import LinearAutoEncoder
plt.style.use('ggplot')
fig, ax = plt.subplots(1, 1)
fig.set_size_inches(12, 8)
fig.suptitle('Pendulum')
ax.set_xlabel('Timestep')
ax.set_ylabel('State')
pend_test_data = PendulumDataset('valid')
states, actions = pend_test_data[583]
states_sim = states.numpy()
ax.plot(states_sim[:, 0], c='r', label='position [true state]', linewidth=2)
ax.plot(states_sim[:, 2], c='b', label='velocity [true state]', linewidth=2)
def trajectory_prediction_lstm():
checkpoints = ['128h_3step_5000_epochs_best.pth']
checkpoints_path = './checkpoints/lstm_auto_encoder/'
model = LSTMAutoEncoder(input_size=3, action_size=1, hidden_size=128, num_layers=1, k_step=1).eval()
for j, checkpoint in enumerate(checkpoints):
checkpoint_path = checkpoints_path + checkpoint
model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')), strict=True)
states_net = torch.zeros(200, 3)
states_net[0] = states[0]
state_t = states[0].view(1, 1, 3)
h_t = torch.zeros(1, 1, 128)
c_t = torch.zeros(1, 1, 128)
# with torch.no_grad():
# for t in range(states.size(0)-1):
# encoded = model.encode(state_t)
# encoded, (h_t, c_t) = model.lstm(encoded, (h_t, c_t))
# transformed = model.transform(encoded, actions[t])
# state_t = model.decode(transformed)
# states_net[t+1] = state_t.squeeze()
with torch.no_grad():
state_hidden_t = model.encode(state_t)
state_hidden_t, (h_t, c_t) = model.lstm(state_hidden_t, (h_t, c_t))
for t in range(states.size(0)-1):
next_state_hiddent_t = model.transform(state_hidden_t, actions[t])
state_t = model.decode(next_state_hiddent_t)
states_net[t+1] = state_t
state_hidden_t = next_state_hiddent_t
# with torch.no_grad():
# for t in range(states.size(0)-1):
# encoded, (h_t, c_t) = model.lstm(state_t, (h_t, c_t))
# transformed = model.transform(encoded, actions[t])
# state_t = model.f_decoder(transformed)
# states_net[t+1] = state_t.squeeze()
# with torch.no_grad():
# state_encoded, _ = model.lstm(states[0].view(1, 1, 3))
# for i in range(states.size(0)-1):
# with torch.no_grad():
# state_encoded = model.transform(state_encoded, actions[i])
# state_decoded = model.f_decoder(state_encoded)
# # print(state_decoded.shape)
# states_net[0, i+1] = state_decoded[:, -1, :]
states_net = states_net.numpy()
ax.plot(states_net[:, 0], '--', label='position - 128h - w/ lstm', linewidth=2)
ax.plot(states_net[:, 2], '--', label='velocity - 128h - w/ lstm', linewidth=2)
def trajectory_prediction_linear():
checkpoints_path = './checkpoints/linear_auto_encoder/'
checkpoints = ['checkpoint_16h_1step_linear.pth']
# checkpoint_path = './checkpoints/lstm_auto_encoder/checkpoint_16h_2step.pth'
model = LinearAutoEncoder(input_size=3, action_size=1, hidden_size=16, bias=True, k_step=1).eval()
for j, checkpoint in enumerate(checkpoints):
checkpoint_path = checkpoints_path + checkpoint
model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')), strict=True)
states_net = torch.zeros(1, 200, 3)
states_net[0, 0] = states[0]
for i in range(states.size(0)-1):
with torch.no_grad():
encoded, decoded = model.forward(states_net[:, :i+1, :])
transformed = model.transform(encoded, actions[:i+1])
transformed_decoded = model.decode(transformed)
transformed_decoded = transformed_decoded[:, -1, :]
states_net[0, i+1] = transformed_decoded
states_net = states_net.view(200, 3).numpy()
ax.plot(states_net[:, 0], '--', label='position - ' + checkpoint, linewidth=2)
ax.plot(states_net[:, 2], '--', label='velocity - ' + checkpoint, linewidth=2)
if __name__ == '__main__':
trajectory_prediction_lstm()
# trajectory_prediction_linear()
plt.legend()
plt.show()