def __call__(self, arr, aux=None):
shape = arr.shape
if self.move is not None:
center = random_center(shape, self.move)
arr_ret = crop(arr, center, self.size)
angle = np.random.randint(4, size=3)
arr_ret = rotation(arr_ret, angle=angle)
axis = np.random.randint(4) - 1
arr_ret = reflection(arr_ret, axis=axis)
# order = np.random.permutation(3) # generate order
# arr_ret = reorder(arr_ret, order)
# zoom_rate = 0.8 + (1.15 - 0.8) * np.random.rand()
# (arr_ret, _) = resize(arr_ret, [1., 1., 1.], new_spacing=[zoom_rate] * 3)
# zoomed_center = arr_ret.shape // 2
# arr_ret = crop_at_zyx_with_dhw(arr_ret, zoomed_center, self.size, 0)
arr_ret = np.expand_dims(arr_ret, axis=-1)
if aux is not None:
aux_ret = crop(aux, center, self.size)
aux_ret = rotation(aux_ret, angle=angle)
aux_ret = reflection(aux_ret, axis=axis)
# aux_ret = reorder(aux_ret, order=order)
# (aux_ret, _) = resize(aux_ret, [1., 1., 1.], new_spacing=[zoom_rate] * 3)
# aux_ret = crop_at_zyx_with_dhw(aux_ret, zoomed_center, self.size, 0)
aux_ret = np.expand_dims(aux_ret, axis=-1)
return arr_ret, aux_ret
return arr_ret
else:
center = np.array(shape) // 2
arr_ret = crop(arr, center, self.size)
arr_ret = np.expand_dims(arr_ret, axis=-1)
if aux is not None:
aux_ret = crop(aux, center, self.size)
aux_ret = np.expand_dims(aux_ret, axis=-1)
return arr_ret, aux_ret
return arr_ret
def draw_ant(draw: ImageDraw.ImageDraw, ant: Ant):
y, x = ant.position
pos = (np.array([x, y]) + .5) * c.C2P
rot = rotation(-ant.direction * TAU / 8)
start = rot @ np.array([0., -c.C2P / 3]) + pos
end = rot @ np.array([0., c.C2P / 3]) + pos
color = (255, 0, 0) if ant.food else (0, 0, 0)
draw.line(start.tolist() + end.tolist(), fill=color, width=2)
def draw_ant(ant: Ant, surface: Surface):
y, x = ant.position
pos = (np.array([x, y]) + 0.5) * c.C2P
rot = rotation(-ant.direction * TAU / 8)
start = (rot @ np.array([0., -c.C2P / 3]) + pos).astype(int)
end = (rot @ np.array([0., c.C2P / 3]) + pos).astype(int)
color = (255, 0, 0) if ant.food else (0, 0, 0)
# pygame.draw.line(surface, color, start.tolist(), end.tolist(), width=1)
draw_arrow(surface, color, end.tolist(), start.tolist(), width=1)
def dead_reckoning(self, points, dt, v, w):
if w == 0:
# going in a straight line.
displacement = dt * v
points[:, 0] -= displacement
else:
angle_along_arc = dt * w
radius_of_curvature = np.abs(v / w)
dx = radius_of_curvature * np.sin(angle_along_arc)
dy = radius_of_curvature * (1 - np.cos(angle_along_arc))
# print("dx:", dx, "dy:", dy)
points[:, 0] -= dx
points[:, 1] -= dy
rotation_matrix = rotation(angle_along_arc)
points[:, :2] = np.matmul(points[:, :2], rotation_matrix)
return points
def motion_cb(self, msg):
new_pose = np.zeros((3, 1))
new_pose[0] = msg.pose.pose.position.x
new_pose[1] = msg.pose.pose.position.y
new_pose[2] = utils.quaternion_to_angle(msg.pose.pose.orientation)
self.state_lock.acquire()
if isinstance(self.last_pose, np.ndarray):
dx = new_pose[0, 0] - self.last_pose[0, 0]
dy = new_pose[1, 0] - self.last_pose[1, 0]
d_theta = new_pose[2, 0] - self.last_pose[2, 0]
# Compute the control from the msg and last_pose
# Apply noise.
# TODO: May want to scale e_theta
# Rotation matrix
rotation = utils.rotation(self.last_pose[2, 0])
dX = np.array((dx, dy)).reshape((2, 1))
dX_prime = np.matmul(rotation.transpose(), dX)
control = np.array((dX_prime[0, 0], dX_prime[1, 0], d_theta))
self.apply_motion_model(self.particles, control)
self.last_pose = new_pose
self.state_lock.release()
def transform(arr):
angle = np.random.randint(4, size=3)
arr_ret = rotation(arr, angle=angle)
axis = np.random.randint(4) - 1
arr_ret = reflection(arr_ret, axis=axis)
return arr_ret
def stage_cost(env, goal, FLAG):
stg_cost = np.full((n, n, 4, 3), np.inf)
for i in range(n):
for j in range(n):
for direc in range(4):
for u in range(3):
if FLAG and what([i, j], env) == 'Door':
stg_cost[i, j, :, :] = 20000
continue
if what([i, j], env) == 'Wall':
stg_cost[i, j, :, :] = 20000
continue
# Move Forward
if u == 0:
# position of next step
pos = np.array([i, j]) + direcs[direc]
cost_fw = 1
if FLAG and what([i, j], env) == 'Door':
cost_dis = 10000
if what([pos[0], pos[1]], env) == 'Wall':
cost_dis = 10000
else:
# distance b/t next step and the goal
cost_dis = (pos[0] - goal[0])**2 + (pos[1] -
goal[1])**2
stg_cost[i, j, direc, u] = cost_fw + cost_dis
# Turn Left
if u == 1:
# the new direction after rotation
angle = (rotation(-90) @ np.transpose(
direcs[direc])).astype(int)
# position that rotate and move one step forward
pos_new = np.array([i, j]) + angle
# default cost
cost_rt = 5
# distance b/t original pos
cost_tr_dis = (i - goal[0])**2 + (j - goal[1])**2
if FLAG and what([pos_new[0], pos_new[1]],
env) == 'Door':
cost_dis = 10000
if what([pos_new[0], pos_new[1]], env) == 'Wall':
cost_fw_dis = 10000
else:
# distance b/t the next step after a turn
cost_fw_dis = ((pos_new[0] - goal[0])**2 +
(pos_new[1] - goal[1])**2)
stg_cost[i, j, direc,
u] = cost_rt + cost_tr_dis + cost_fw_dis
# Turn Right
if u == 2:
# the new direction after rotation
angle = (rotation(90) @ np.transpose(
direcs[direc])).astype(int)
# position that rotate and move one step forward
pos_new = np.array([i, j]) + angle
cost_rt = 5
cost_tr_dis = (i - goal[0])**2 + (j - goal[1])**2
if FLAG and what([pos_new[0], pos_new[1]],
env) == 'Door':
cost_dis = 10000
if what([pos_new[0], pos_new[1]], env) == 'Wall':
cost_fw_dis = 10000
else:
# distance b/t the next step after a turn
cost_fw_dis = ((pos_new[0] - goal[0])**2 +
(pos_new[1] - goal[1])**2)
stg_cost[i, j, direc,
u] = cost_rt + cost_tr_dis + cost_fw_dis
#print("tr cost:", stg_cost[i,j,direc,u])
return stg_cost
def update(v_T, stg_cost, T):
v_t = np.full((n, n, 4), np.inf)
policy = np.empty([n, n, 4, T])
for t in range(T - 1, -1, -1):
#print(t)
Q_t = np.full((n, n, 4, 3), np.inf)
# use terminal cost for the first step
if t == T - 1:
v = copy.deepcopy(v_T)
else:
v = copy.deepcopy(v_t)
# calculate total cost
for i in range(n):
for j in range(n):
for direc in range(4):
for u in range(3):
if u == 0:
pos = np.array([i, j]) + direcs[direc]
if pos[0] < n and pos[0] >= 0 and pos[
1] < n and pos[1] >= 0:
v_val = v[pos[0], pos[1], direc]
else:
v_val = 10000
if u == 1:
angle = (rotation(-90) @ np.transpose(
direcs[direc])).astype(int)
v_val = v[i, j, dir2num(angle)]
if u == 2:
angle = (rotation(90) @ np.transpose(
direcs[direc])).astype(int)
v_val = v[i, j, dir2num(angle)]
# given a state, stage cost + the cost of its next step
Q_t[i, j, direc, u] = stg_cost[i, j, direc, u] + v_val
#find optimal cost of each states
for i in range(n):
for j in range(n):
for direc in range(4):
# find the index of the minima
motion = np.where(
Q_t[i, j, direc, :] == np.min(Q_t[i, j,
direc, :]))[0][0]
opt_cost = Q_t[i, j, direc, motion]
v_t[i, j, direc] = opt_cost
policy[i, j, direc, t] = motion
policy = policy.astype(int)
# report collect
key_s1_.append(v_t[2, 1, 1])
key_s2_.append(v_t[1, 2, 2])
#key_s3_.append(v_t[1,2,2])
#key_s4_.append(v_t[2,6,3])
goal_s1_.append(v_t[3, 1, 0])
#goal_s2_.append(v_t[4,3,2])
door_s1_.append(v_t[2, 1, 0])
if (v_t == v).all():
break
return v_t, policy
# Iterator
dataiter = iter(train_loader_match)
if args.gs:
kernel_size = 3
pad = 2
sigma = 1
kernel = get_gaussian_kernel(kernel_size=kernel_size, pad=pad,
sigma=sigma).cuda()
criterion_kl = nn.KLDivLoss(size_average=False)
for epoch in range(args.epochs):
running_loss = 0
for i, (imgs, _) in enumerate(train_loader):
img = imgs[0].to(device)
img_rot = rotation(img)[0]
img_aug = imgs[1].to(device)
try:
img_match = next(dataiter)[0]
except StopIteration:
dataiter = iter(train_loader_match)
img_match = next(dataiter)[0]
img_match = img_match.to(device)
netG.train()
optimG.zero_grad()
# Unconstrained Adversaries
adv = netG(img)
adv_rot = netG(img_rot)