def train(self, in_img, p, q, theta, backprop=True):
self.metric.reset_states()
with tf.GradientTape() as tape:
output = self.forward(in_img, p, softmax=False)
itheta = theta / (2 * np.pi / self.n_rotations)
itheta = np.int32(np.round(itheta)) % self.n_rotations
label_size = in_img.shape[:2] + (self.n_rotations, )
label = np.zeros(label_size)
label[q[0], q[1], itheta] = 1
# Get per-pixel sampling loss.
sampling = True # Sampling negatives seems to converge faster.
if sampling:
num_samples = 100
inegative = utils.sample_distribution(1 - label, num_samples)
inegative = [
np.ravel_multi_index(i, label.shape) for i in inegative
]
ipositive = np.ravel_multi_index([q[0], q[1], itheta],
label.shape)
output = tf.reshape(output, (-1, 2))
output_samples = ()
for i in inegative:
output_samples += (tf.reshape(output[i, :], (1, 2)), )
output_samples += (tf.reshape(output[ipositive, :], (1, 2)), )
output = tf.concat(output_samples, axis=0)
label = np.int32([0] * num_samples + [1])[Ellipsis, None]
label = np.hstack((1 - label, label))
weights = np.ones(label.shape[0])
weights[:num_samples] = 1. / num_samples
weights = weights / np.sum(weights)
else:
ipositive = np.ravel_multi_index([q[0], q[1], itheta],
label.shape)
output = tf.reshape(output, (-1, 2))
label = np.int32(
np.reshape(label, (int(np.prod(label.shape)), 1)))
label = np.hstack((1 - label, label))
weights = np.ones(label.shape[0]) * 0.0025 # Magic constant.
weights[ipositive] = 1
label = tf.convert_to_tensor(label, dtype=tf.int32)
weights = tf.convert_to_tensor(weights, dtype=tf.float32)
loss = tf.nn.softmax_cross_entropy_with_logits(label, output)
loss = tf.reduce_mean(loss * weights)
train_vars = self.model.trainable_variables
if backprop:
grad = tape.gradient(loss, train_vars)
self.optim.apply_gradients(zip(grad, train_vars))
self.metric(loss)
self.iters += 1
return np.float32(loss)
def train(self, input_image, p, q, theta):
"""Train function."""
self.metric.reset_states()
with tf.GradientTape() as tape:
output = self.forward(input_image)
p_descriptor = output[0, p[0], p[1], :]
itheta = theta / (2 * np.pi / self.num_rotations)
itheta = np.int32(np.round(itheta)) % self.num_rotations
q_descriptor = output[itheta, q[0], q[1], :]
# Positives.
positive_distances = tf.linalg.norm(p_descriptor - q_descriptor)
positive_distances = tf.reshape(positive_distances, (1, ))
positive_labels = tf.constant([1], dtype=tf.int32)
positive_loss = tfa_losses.contrastive_loss(
positive_labels, positive_distances)
# Negatives.
num_samples = 100
sample_map = np.zeros(input_image.shape[:2] +
(self.num_rotations, ))
sample_map[p[0], p[1], 0] = 1
sample_map[q[0], q[1], itheta] = 1
inegative = utils.sample_distribution(1 - sample_map, num_samples)
negative_distances = ()
negative_labels = ()
for i in range(num_samples):
descriptor = output[inegative[i, 2], inegative[i, 0],
inegative[i, 1], :]
distance = tf.linalg.norm(p_descriptor - descriptor)
distance = tf.reshape(distance, (1, ))
negative_distances += (distance, )
negative_labels += (tf.constant([0], dtype=tf.int32), )
negative_distances = tf.concat(negative_distances, axis=0)
negative_labels = tf.concat(negative_labels, axis=0)
negative_loss = tfa_losses.contrastive_loss(
negative_labels, negative_distances)
negative_loss = tf.reduce_mean(negative_loss)
loss = tf.reduce_mean(positive_loss) + tf.reduce_mean(
negative_loss)
# Backpropagate.
grad = tape.gradient(loss, self.model.trainable_variables)
self.optim.apply_gradients(zip(grad, self.model.trainable_variables))
self.metric(loss)
return np.float32(loss)
def get_random_pose(self, env, obj_size):
"""Get random collision-free object pose within workspace bounds."""
# Get erosion size of object in pixels.
max_size = np.sqrt(obj_size[0] ** 2 + obj_size[1] ** 2)
erode_size = int(np.round(max_size / self.pix_size))
_, hmap, obj_mask = self.get_true_image(env)
# Randomly sample an object pose within free-space pixels.
free = np.ones(obj_mask.shape, dtype=np.uint8)
for obj_ids in env.obj_ids.values():
for obj_id in obj_ids:
free[obj_mask == obj_id] = 0
free[0, :], free[:, 0], free[-1, :], free[:, -1] = 0, 0, 0, 0
free = cv2.erode(free, np.ones((erode_size, erode_size), np.uint8))
if np.sum(free) == 0:
return None, None
pix = utils.sample_distribution(np.float32(free))
pos = utils.pix_to_xyz(pix, hmap, self.bounds, self.pix_size)
pos = (pos[0], pos[1], obj_size[2] / 2)
theta = np.random.rand() * 2 * np.pi
rot = utils.eulerXYZ_to_quatXYZW((0, 0, theta))
return pos, rot
def act(obs, info): # pylint: disable=unused-argument
"""Calculate action."""
# Oracle uses perfect RGB-D orthographic images and segmentation masks.
_, hmap, obj_mask = self.get_true_image(env)
# Unpack next goal step.
objs, matches, targs, replace, rotations, _, _, _ = self.goals[0]
# Match objects to targets without replacement.
if not replace:
# Modify a copy of the match matrix.
matches = matches.copy()
# Ignore already matched objects.
for i in range(len(objs)):
object_id, (symmetry, _) = objs[i]
pose = p.getBasePositionAndOrientation(object_id)
targets_i = np.argwhere(matches[i, :]).reshape(-1)
for j in targets_i:
# SAY: check whether the object arrives its target.
if self.is_match(pose, targs[j], symmetry):
matches[i, :] = 0
matches[:, j] = 0
# Get objects to be picked (prioritize farthest from nearest neighbor).
nn_dists = []
nn_targets = []
for i in range(len(objs)):
object_id, (symmetry, _) = objs[i]
xyz, _ = p.getBasePositionAndOrientation(object_id)
targets_i = np.argwhere(matches[i, :]).reshape(-1)
if len(targets_i) > 0: # pylint: disable=g-explicit-length-test
targets_xyz = np.float32([targs[j][0] for j in targets_i])
dists = np.linalg.norm(
targets_xyz - np.float32(xyz).reshape(1, 3), axis=1)
nn = np.argmin(dists)
nn_dists.append(dists[nn])
nn_targets.append(targets_i[nn])
# Handle ignored objects.
else:
nn_dists.append(0)
nn_targets.append(-1)
order = np.argsort(nn_dists)[::-1]
# SAY: matched objects may be the ones that have been at the target location.
# Filter out matched objects.
order = [i for i in order if nn_dists[i] > 0]
pick_mask = None
for pick_i in order:
pick_mask = np.uint8(obj_mask == objs[pick_i][0])
# Erode to avoid picking on edges.
# pick_mask = cv2.erode(pick_mask, np.ones((3, 3), np.uint8))
if np.sum(pick_mask) > 0:
break
# Trigger task reset if no object is visible.
if pick_mask is None or np.sum(pick_mask) == 0:
self.goals = []
print('Object for pick is not visible. Skipping demonstration.')
return
# Get picking pose.
pick_prob = np.float32(pick_mask)
pick_pix = utils.sample_distribution(pick_prob)
# For "deterministic" demonstrations on insertion-easy, use this:
# pick_pix = (160,80)
pick_pos = utils.pix_to_xyz(pick_pix, hmap,
self.bounds, self.pix_size)
pick_pose = (np.asarray(pick_pos), np.asarray((0, 0, 0, 1)))
# Get placing pose.
targ_pose = targs[nn_targets[pick_i]] # pylint: disable=undefined-loop-variable
obj_pose = p.getBasePositionAndOrientation(objs[pick_i][0]) # pylint: disable=undefined-loop-variable
if not self.sixdof:
obj_euler = utils.quatXYZW_to_eulerXYZ(obj_pose[1])
obj_quat = utils.eulerXYZ_to_quatXYZW((0, 0, obj_euler[2]))
obj_pose = (obj_pose[0], obj_quat)
world_to_pick = utils.invert(pick_pose)
obj_to_pick = utils.multiply(world_to_pick, obj_pose)
pick_to_obj = utils.invert(obj_to_pick)
place_pose = utils.multiply(targ_pose, pick_to_obj)
# Rotate end effector?
if not rotations:
place_pose = (place_pose[0], (0, 0, 0, 1))
place_pose = (np.asarray(place_pose[0]), np.asarray(place_pose[1]))
return {'pose0': pick_pose, 'pose1': place_pose}