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}