def random_pose_6dof(self, env, object_size): """Get random collision-free pose in workspace bounds for object.""" plane_id = 1 max_size = np.linalg.norm(object_size[0:2]) erode_size = int(np.round(max_size / self.pixel_size)) _, heightmap, object_mask = self.get_object_masks(env) # Sample freespace regions in workspace. mask = np.uint8(object_mask == plane_id) mask[0, :], mask[:, 0], mask[-1, :], mask[:, -1] = 0, 0, 0, 0 mask = cv2.erode(mask, np.ones((erode_size, erode_size), np.uint8)) if np.sum(mask) == 0: return pixel = utils.sample_distribution(np.float32(mask)) position = utils.pixel_to_position(pixel, heightmap, self.bounds, self.pixel_size) z_above_table = (np.random.rand(1)[0] / 10) + 0.03 position = (position[0], position[1], object_size[2] / 2 + z_above_table) roll = (np.random.rand() - 0.5) * 0.5 * np.pi pitch = (np.random.rand() - 0.5) * 0.5 * np.pi yaw = np.random.rand() * 2 * np.pi rotation = utils.get_pybullet_quaternion_from_rot((roll, pitch, yaw)) print(position, rotation) return position, rotation
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 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: 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] # 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 = (pick_pos, (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)) return {'pose0': pick_pose, 'pose1': place_pose}
def train(self, in_img, p, q, theta): """Transport pixel p to pixel q. Args: input: depth_image: p: pixel (y, x) q: pixel (y, x) Returns: A `Tensor`. Has the same type as `input`. Daniel: the `in_img` will include the color and depth. Much is similar to the attention model if we're not using the per-pixel loss: (a) forward pass, (b) get angle discretizations [though we set only 1 rotation for the picking model], (c) make the label consider rotations in the last axis, but only provide the label to one single (pixel,rotation) combination, (d) follow same exact steps for the non-per pixel loss otherwise. The output reshaping to (1, ...) is done in the attention model forward pass, but not in the transport forward pass. Note the `1` meaning a batch size of 1. """ self.metric.reset_states() with tf.GradientTape() as tape: output = self.forward(in_img, p, apply_softmax=False) itheta = theta / (2 * np.pi / self.num_rotations) itheta = np.int32(np.round(itheta)) % self.num_rotations label_size = in_img.shape[:2] + (self.num_rotations, ) label = np.zeros(label_size) label[q[0], q[1], itheta] = 1 if self.per_pixel_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])[..., 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) else: label = label.reshape(1, np.prod(label.shape)) label = tf.convert_to_tensor(label, dtype=tf.float32) output = tf.reshape(output, (1, np.prod(output.shape))) loss = tf.nn.softmax_cross_entropy_with_logits(label, output) loss = tf.reduce_mean(loss) 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 train(self, in_img, p, q, theta, z=None, roll=None, pitch=None, validate=False): """Transport pixel p to pixel q. Args: in_img: p: pixel (y, x) q: pixel (y, x) theta: z: roll: pitch: validate: Returns: A `Tensor`. Has the same type as `input`. """ visualize_input = False if visualize_input and self.six_dof: # only supported for six dof model self.visualize_train_input(in_img, p, q, theta, z, roll, pitch) self.metric.reset_states() if self.six_dof: self.z_metric.reset_states() self.roll_metric.reset_states() self.pitch_metric.reset_states() with tf.GradientTape() as tape: output = self.forward(in_img, p, apply_softmax=False, theta=theta) if self.six_dof: z_label, roll_label, pitch_label = z, roll, pitch output, z_tensor, roll_tensor, pitch_tensor = output itheta = theta / (2 * np.pi / self.num_rotations) itheta = np.int32(np.round(itheta)) % self.num_rotations label_size = in_img.shape[:2] + (self.num_rotations, ) label = np.zeros(label_size) label[q[0], q[1], itheta] = 1 if self.per_pixel_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) transport_loss = loss elif not self.six_dof: label = label.reshape(1, np.prod(label.shape)) label = tf.convert_to_tensor(label, dtype=tf.float32) output = tf.reshape(output, (1, np.prod(output.shape))) loss = tf.nn.softmax_cross_entropy_with_logits(label, output) loss = tf.reduce_mean(loss) transport_loss = loss if self.six_dof: # Use a window for regression, rather than only exact u_window = 7 v_window = 7 theta_window = 1 u_min = max(q[0] - u_window, 0) u_max = min(q[0] + u_window + 1, z_tensor.shape[1]) v_min = max(q[1] - v_window, 0) v_max = min(q[1] + v_window + 1, z_tensor.shape[2]) theta_min = max(itheta - theta_window, 0) theta_max = min(itheta + theta_window + 1, z_tensor.shape[3]) z_est_at_xytheta = z_tensor[0, u_min:u_max, v_min:v_max, theta_min:theta_max] roll_est_at_xytheta = roll_tensor[0, u_min:u_max, v_min:v_max, theta_min:theta_max] pitch_est_at_xytheta = pitch_tensor[0, u_min:u_max, v_min:v_max, theta_min:theta_max] z_est_at_xytheta = tf.reshape(z_est_at_xytheta, (-1, 1)) roll_est_at_xytheta = tf.reshape(roll_est_at_xytheta, (-1, 1)) pitch_est_at_xytheta = tf.reshape(pitch_est_at_xytheta, (-1, 1)) z_est_at_xytheta = self.z_regressor(z_est_at_xytheta) roll_est_at_xytheta = self.roll_regressor(roll_est_at_xytheta) pitch_est_at_xytheta = self.pitch_regressor( pitch_est_at_xytheta) z_weight = 10.0 roll_weight = 10.0 pitch_weight = 10.0 z_label = tf.convert_to_tensor(z_label)[None, Ellipsis] roll_label = tf.convert_to_tensor(roll_label)[None, Ellipsis] pitch_label = tf.convert_to_tensor(pitch_label)[None, Ellipsis] z_loss = z_weight * self.regress_loss(z_label, z_est_at_xytheta) roll_loss = roll_weight * self.regress_loss( roll_label, roll_est_at_xytheta) pitch_loss = pitch_weight * self.regress_loss( pitch_label, pitch_est_at_xytheta) loss = z_loss + roll_loss + pitch_loss if self.six_dof: train_vars = self.model.trainable_variables + \ self.z_regressor.trainable_variables + \ self.roll_regressor.trainable_variables + \ self.pitch_regressor.trainable_variables else: train_vars = self.model.trainable_variables if not validate: grad = tape.gradient(loss, train_vars) self.optim.apply_gradients(zip(grad, train_vars)) if not self.six_dof: self.metric(transport_loss) if self.six_dof: self.z_metric(z_loss) self.roll_metric(roll_loss) self.pitch_metric(pitch_loss) return np.float32(loss)