-
Notifications
You must be signed in to change notification settings - Fork 2
/
test_module.py
373 lines (276 loc) · 13.7 KB
/
test_module.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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
import numpy as np
from matplotlib import pyplot as plt
import torch
import torch.nn.functional as F
from dataloader.constants import *
from dataloader.gnn_setup import *
from graphs.losses.cross_entropy import CrossEntropyLoss
print('Parameters used')
def generate_data(data_size):
feat_list, adj_list, label_list = [], [], []
grid_list = []
robot_pos_list = []
for _ in range(data_size):
grid = get_reward_grid(height=HEIGHT, width=WIDTH, reward_thresh=REWARD_THRESH)
robot_pos, adj_mat = get_initial_pose(grid, comm_range=COMM_RANGE)
cent_act, cent_rwd = centralized_greedy_action_finder(grid, robot_pos, fov=FOV)
rand_act, rand_rwd = random_action_finder(grid, robot_pos, 1000)
if cent_rwd > rand_rwd:
action_vec = cent_act
else:
action_vec = rand_act
feat_vec = get_features(grid, robot_pos, fov=FOV, step=STEP, target_feat_size=NUM_TGT_FEAT, robot_feat_size=NUM_ROB_FEAT)
feat_list.append(feat_vec)
adj_list.append(adj_mat)
action_one_hot = np.zeros((NUM_ROBOT, len(DIR_LIST)), dtype=np.uint8)
action_one_hot[np.arange(NUM_ROBOT), action_vec] = 1
label_list.append(action_one_hot)
grid_list.append(grid.copy())
robot_pos_list.append(robot_pos.copy())
return [np.array(feat_list), np.array(adj_list), np.array(label_list), grid_list, robot_pos_list]
def calculate_test_reward(grid, robot_pos, action_list):
"""
Function to calculate the reward calculated by all the robots based on an action vector.
For this we first update locations of each robot, then create a mask which has 1s only around the new robots locations (square of side (2*FOV+1) for each robot)
Parameters
----------
grid: 2D grid containing rewards
robot_pos: Current position for each robot on the grid (NUM_ROBOTx2 size vector)
action_list: List of action for each robot
Returns
-------
total_reward: Total reward calculated by the robots using action_list (the action vector)
"""
# Convert the integer actions to 2D vector of location differences using DIR_DICT dictionary
act = np.array([DIR_DICT[k] for k in action_list])
# Calcuate new locations for each robot
new_pos = robot_pos + act
# Make sure that the new locatiosn are within the grid
new_pos = new_pos.clip(min=0, max=GRID_SIZE-1)
# Initialize a mask of same shape as grid
mask = np.zeros(grid.shape, dtype=int)
# iterate over each robot position
for c_pos, n_pos in zip(robot_pos, new_pos):
# Set the values to 1 in the mask at each robot's fov
# also make sure that the indices do not go out of grid
# Calculate the bounding box ranges for the box generated by robot moving from the current location (c_pos) to new location (n_pos)
# This box has a padding of size FOV on each size
r_lim_lef = max(0, min(c_pos[0]-FOV, n_pos[0]-FOV))
c_lim_top = max(0, min(c_pos[1]-FOV, n_pos[1]-FOV))
r_lim_rgt = min(max(c_pos[0]+FOV+1, n_pos[0]+FOV+1), GRID_SIZE)
c_lim_bot = min(max(c_pos[1]+FOV+1, n_pos[1]+FOV+1), GRID_SIZE)
# Set the locations withing mask (i.e. witing robot's vision when it moved) to 1
mask[r_lim_lef:r_lim_rgt, c_lim_top:c_lim_bot] = 1
# Find total reward as number of 1s in the masked grid
total_reward = np.sum(grid * mask)
return total_reward, mask
def plot_function(grid, robot_pos, gt_act, predict_act, random_act, greedy=True):
gt_rwd, gt_mask = calculate_test_reward(grid, robot_pos, gt_act)
pred_rwd, pred_mask = calculate_test_reward(grid, robot_pos, predict_act)
rndm_rwd, rndm_mask = calculate_test_reward(grid, robot_pos, random_act)
num_row = 1
num_col = 3
if greedy:
cent_act, cent_rwd = centralized_greedy_action_finder(grid, robot_pos, fov=FOV)
cent_rwd, cent_mask = calculate_test_reward(grid, robot_pos, cent_act)
num_row = 2
num_col = 2
if greedy:
plt.figure(figsize=(7,8))
else:
plt.figure(figsize=(12,6));
plt.subplot(num_row,num_col,1);
plt.imshow(gt_mask);
plt.title(f'GT: {gt_rwd}', fontsize= 20)
plt.subplot(num_row,num_col,2);
plt.imshow(pred_mask);
plt.title(f'Prediction: {pred_rwd}', fontsize= 20)
plt.subplot(num_row,num_col,3);
plt.imshow(rndm_mask);
plt.title(f'Random: {rndm_rwd}', fontsize= 20)
if greedy:
plt.subplot(num_row,num_col, 4)
plt.imshow(cent_mask);
plt.title(f'C. Greedy: {cent_rwd}', fontsize=20)
# plt.suptitle('Reward calculation', fontsize= 20);
def get_accuracy(config, agent, data_loader):
# data_loader = agent.data_loader.test_loader
gt_list_long, pred_list_long = [], []
agent.model.eval();
for batch_idx, (batch_input, batch_GSO, batch_target) in enumerate(data_loader):
inputGPU = batch_input.to(config.device)
gsoGPU = batch_GSO.to(config.device)
# gsoGPU = gsoGPU.unsqueeze(0)
targetGPU = batch_target.to(config.device)
# Should not transpose if using flattened batch
# batch_targetGPU = targetGPU.permute(1, 0, 2)
batch_targetGPU = targetGPU
agent.optimizer.zero_grad()
# loss
loss_validStep = 0
# model
agent.model.addGSO(gsoGPU)
predict = agent.model(inputGPU)
gt_list_long.append(targetGPU.detach().cpu().numpy())
pred_list_long.append(predict.detach().cpu().numpy())
# pred_list_long.append(np.array([p.detach().cpu().numpy() for p in predict]).transpose(1,0,2))
np.concatenate(gt_list_long, axis=0).shape, np.concatenate(pred_list_long, axis=0).shape
gt_idxs = np.concatenate(gt_list_long, axis=0).argmax(axis=2)
pred_idxs = np.concatenate(pred_list_long, axis=0).argmax(axis=2)
accuracy = (gt_idxs == pred_idxs).sum()/(len(gt_idxs)*config.num_agents)
return accuracy
def main(config, agent, num_example=5):
temp_list = generate_data(num_example)
print(config.tgt_feat, config.rbt_feat)
numFeature = (config.tgt_feat + config.rbt_feat )
featlist, adjlist, tgtlist = [], [], []
feat, adj, tgt, _, _ = temp_list
feat_reshaped = feat[:,:,:numFeature,:].reshape(feat.shape[0], feat.shape[1], numFeature*2)
featlist.append(feat_reshaped)
adjlist.append(adj)
tgtlist.append(tgt)
features_tensor = torch.FloatTensor(np.concatenate(featlist, axis=0))
adj_mat_tensor = torch.FloatTensor(np.concatenate(adjlist, axis=0))
targets_tensor = torch.LongTensor(np.concatenate(tgtlist, axis=0))
inputGPU = features_tensor.to(config.device)
gsoGPU = adj_mat_tensor.to(config.device)
# gsoGPU = gsoGPU.unsqueeze(0)
targetGPU = targets_tensor.to(config.device)
# Should not transpose if using flattened batch
# batch_targetGPU = targetGPU.permute(1, 0, 2)
batch_targetGPU = targetGPU
agent.model.eval();
agent.optimizer.zero_grad()
agent.model.addGSO(gsoGPU)
predict = agent.model(inputGPU)
predict_np = predict.detach().cpu().numpy()
# predict_np = np.array([p.detach().cpu().numpy() for p in predict])
predict_ind = predict_np.argmax(axis=2)
_, _, gt_act_list, grid_list, robot_pos_list = temp_list
gt_act_list = gt_act_list.argmax(axis=2)
random_acts = np.random.randint(low=0,high=len(DIR_DICT.keys()), size=(len(gt_act_list), config.num_agents))
return grid_list, robot_pos_list, gt_act_list, predict_ind, random_acts
def get_acc_n_loss(config, agent, data_loader):
# data_loader = agent.data_loader.test_loader
gt_list_long, pred_list_long = [], []
agent.model.eval();
loss_fn = CrossEntropyLoss()
loss_fn = loss_fn.to(config.device)
agent.model.eval()
with torch.no_grad():
for batch_idx, (batch_input, batch_GSO, batch_target) in enumerate(data_loader):
inputGPU = batch_input.to(config.device)
gsoGPU = batch_GSO.to(config.device)
# gsoGPU = gsoGPU.unsqueeze(0)
targetGPU = batch_target.to(config.device)
# Should not transpose if flattening the batch
# batch_targetGPU = targetGPU.permute(1, 0, 2)
batch_targetGPU = targetGPU
agent.optimizer.zero_grad()
# print('Data shapes: ', inputGPU.shape, gsoGPU.shape)
# model
agent.model.addGSO(gsoGPU)
predict = agent.model(inputGPU)
gt_list_long.append(targetGPU.detach().cpu().numpy())
pred_list_long.append(predict.detach().cpu().numpy())
# pred_list_long.append(np.array([p.detach().cpu().numpy() for p in predict]).transpose(1,0,2))
loss_val = 0
num_act = batch_targetGPU.shape[-1]
assert predict.shape == batch_targetGPU.shape
loss_val = loss_val + loss_fn(predict.reshape(-1,num_act), torch.max(batch_targetGPU.reshape(-1,num_act), 1)[1])
'''
# Not needed for flattened inputs
# print(NUM_ROBOT, len(predict))
for id_agent in range(NUM_ROBOT):
# for output, target in zip(predict, target):
batch_predict_currentAgent = predict[id_agent][:]
batch_target_currentAgent = batch_targetGPU[id_agent][:][:]
loss_val = loss_val + loss_fn(batch_predict_currentAgent, torch.max(batch_target_currentAgent, 1)[1])
'''
print(f'Num agents: {config.num_agents}')
print(np.concatenate(gt_list_long, axis=0).shape, np.concatenate(pred_list_long, axis=0).shape)
gt_idxs = np.concatenate(gt_list_long, axis=0).argmax(axis=2)
pred_idxs = np.concatenate(pred_list_long, axis=0).argmax(axis=2)
accuracy = (gt_idxs == pred_idxs).sum()/(len(gt_idxs)*config.num_agents)
loss_val = loss_val #/config.num_agents
log_loss = loss_val.item()
return accuracy, log_loss
def get_stoc_acc(config, agent, data_loader):
# data_loader = agent.data_loader.test_loader
gt_list_long, pred_list_long = [], []
agent.model.eval();
with torch.no_grad():
for batch_idx, (batch_input, batch_GSO, batch_target) in enumerate(data_loader):
inputGPU = batch_input.to(config.device)
gsoGPU = batch_GSO.to(config.device)
# gsoGPU = gsoGPU.unsqueeze(0)
targetGPU = batch_target.to(config.device)
# Should not transpose if flattening the batch
# batch_targetGPU = targetGPU.permute(1, 0, 2)
batch_targetGPU = targetGPU
agent.optimizer.zero_grad()
# model
agent.model.addGSO(gsoGPU)
predict = agent.model(inputGPU)
gt_list_long.append(targetGPU.detach().cpu().numpy())
pred_list_long.append(F.softmax(predict, dim=2).detach().cpu().numpy() )
# pred_list_long.append(np.array([F.softmax(p, dim=1).detach().cpu().numpy() for p in predict]).transpose(1,0,2))
np.concatenate(gt_list_long, axis=0).shape, np.concatenate(pred_list_long, axis=0).shape
gt_idxs = np.concatenate(gt_list_long, axis=0)
gt_idxs = gt_idxs.reshape(-1,5).argmax(axis=1)
# print(gt_idxs[0:4])
# gt_idxs = gt_idxs.argmax(axis=2)
pred_idxs = np.concatenate(pred_list_long, axis=0)
pred_idxs = pred_idxs.reshape(-1,5)
# print(pred_idxs[0:4])
pred_idxs = np.array([np.random.choice(np.arange(5), p=logit) for logit in pred_idxs])
# print(pred_idxs.shape)
# pred_idxs = pred_idxs.argmax(axis=2)
accuracy = (gt_idxs == pred_idxs).sum()/(len(gt_idxs))
return accuracy
def convert_to_action(pred_list):
preds = pred_list
'''
# not needed for flattened batch
preds = torch.zeros(pred_list[0].shape[0], len(pred_list), pred_list[0].shape[1])
for itr in range(len(pred_list)):
preds[:,itr,:] = pred_list[itr]
'''
preds = F.softmax(preds, dim=2)
action_ids = torch.zeros((preds.shape[:2]), dtype=torch.long)
for itr in range(preds.shape[1]):
action_ids[:,itr] = torch.multinomial(preds[:,itr,:], 1)[:,0]
return action_ids
def get_stoc_acc2(config, agent, data_loader):
# data_loader = agent.data_loader.test_loader
gt_list_long, pred_list_long = [], []
agent.model.eval();
with torch.no_grad():
for batch_idx, (batch_input, batch_GSO, batch_target) in enumerate(data_loader):
inputGPU = batch_input.to(config.device)
gsoGPU = batch_GSO.to(config.device)
# gsoGPU = gsoGPU.unsqueeze(0)
targetGPU = batch_target.to(config.device)
# Should not transpose if using flattened batch
# batch_targetGPU = targetGPU.permute(1, 0, 2)
batch_targetGPU = targetGPU
agent.optimizer.zero_grad()
# model
agent.model.addGSO(gsoGPU)
predict = agent.model(inputGPU)
# print(len(predict), predict[0].shape)
probs = convert_to_action(predict)
gt_list_long.append(targetGPU.detach().cpu().numpy())
pred_list_long.append(probs.detach().cpu().numpy().astype(int))
# print(np.concatenate(pred_list_long).shape)
# return 0
# np.concatenate(gt_list_long, axis=0).shape, np.concatenate(pred_list_long, axis=0).shape
gt_idxs = np.concatenate(gt_list_long, axis=0)
gt_idxs = gt_idxs.reshape(-1,5).argmax(axis=1)
# print(gt_idxs[0:4])
# gt_idxs = gt_idxs.argmax(axis=2)
pred_idxs = np.concatenate(pred_list_long, axis=0).astype(int)
pred_idxs = pred_idxs.reshape(-1,)
# print(pred_idxs[0:4])
accuracy = (gt_idxs == pred_idxs).sum()/(len(gt_idxs))
return accuracy