def reward(self): """Get delta rewards for current timestep. Returns: A tuple consisting of the scalar (delta) reward, plus `extras` dict which has extra task-dependent info from the process of computing rewards that gives us finer-grained details. Use `extras` for further data analysis. """ reward, info = 0, {} # Unpack next goal step. objs, matches, targs, _, _, metric, params, max_reward = self.goals[0] # Evaluate by matching object poses. if metric == 'pose': step_reward = 0 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: target_pose = targs[j] if self.is_match(pose, target_pose, symmetry): step_reward += max_reward / len(objs) break # Evaluate by measuring object intersection with zone. elif metric == 'zone': zone_pts, total_pts = 0, 0 obj_pts, zones = params for zone_pose, zone_size in zones: # Count valid points in zone. for obj_id in obj_pts: pts = obj_pts[obj_id] obj_pose = p.getBasePositionAndOrientation(obj_id) world_to_zone = utils.invert(zone_pose) obj_to_zone = utils.multiply(world_to_zone, obj_pose) pts = np.float32(utils.apply(obj_to_zone, pts)) if len(zone_size) > 1: valid_pts = np.logical_and.reduce([ pts[0, :] > -zone_size[0] / 2, pts[0, :] < zone_size[0] / 2, pts[1, :] > -zone_size[1] / 2, pts[1, :] < zone_size[1] / 2, pts[2, :] < self.bounds[2, 1]]) zone_pts += np.sum(np.float32(valid_pts)) total_pts += pts.shape[1] step_reward = max_reward * (zone_pts / total_pts) # Get cumulative rewards and return delta. reward = self.progress + step_reward - self._rewards self._rewards = self.progress + step_reward # Move to next goal step if current goal step is complete. if np.abs(max_reward - step_reward) < 0.01: self.progress += max_reward # Update task progress. self.goals.pop(0) return reward, info
def reset(self, env): self.num_steps = 1 self.goal = {'places': {}, 'steps': []} # Generate randomly shaped box. box_size = self.random_size(0.05, 0.15, 0.05, 0.15, 0.01, 0.06) # Add corner. dimx = (box_size[0] / 2 - 0.025 + 0.0025, box_size[0] / 2 + 0.0025) dimy = (box_size[1] / 2 + 0.0025, box_size[1] / 2 - 0.025 + 0.0025) corner_template = 'assets/corner/corner-template.urdf' replace = {'DIMX': dimx, 'DIMY': dimy} corner_urdf = self.fill_template(corner_template, replace) corner_size = (box_size[0], box_size[1], 0) corner_pose = self.random_pose(env, corner_size) env.add_object(corner_urdf, corner_pose, fixed=True) os.remove(corner_urdf) # Add possible placing poses. theta = utils.get_rot_from_pybullet_quaternion(corner_pose[1])[2] flipped_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, theta + np.pi)) flipped_pose = (corner_pose[0], flipped_rotation) alt_x = (box_size[0] / 2) - (box_size[1] / 2) alt_y = (box_size[1] / 2) - (box_size[0] / 2) alt_position = (alt_x, alt_y, 0) alt_rotation0 = utils.get_pybullet_quaternion_from_rot( (0, 0, np.pi / 2)) alt_rotation1 = utils.get_pybullet_quaternion_from_rot( (0, 0, 3 * np.pi / 2)) alt_pose0 = utils.multiply(corner_pose, (alt_position, alt_rotation0)) alt_pose1 = utils.multiply(corner_pose, (alt_position, alt_rotation1)) self.goal['places'] = { 0: corner_pose, 1: flipped_pose, 2: alt_pose0, 3: alt_pose1 } # Add box. box_template = 'assets/box/box-template.urdf' box_urdf = self.fill_template(box_template, {'DIM': box_size}) box_pose = self.random_pose(env, box_size) box_id = env.add_object(box_urdf, box_pose) os.remove(box_urdf) self.color_random_brown(box_id) self.goal['steps'].append({box_id: (2 * np.pi, [0, 1, 2, 3])})
def reset(self, env): # Add pallet. self.zone_size = (0.3, 0.25, 0.25) zone_urdf = 'assets/pallet/pallet.urdf' rotation = utils.get_pybullet_quaternion_from_rot((0, 0, 0)) self.zone_pose = ((0.5, 0.25, 0.02), rotation) env.add_object(zone_urdf, self.zone_pose, fixed=True) # Add stack of boxes on pallet. margin = 0.01 self.object_points = {} stack_size = (0.19, 0.19, 0.19) box_template = 'assets/box/box-template.urdf' stack_dim = np.random.randint(low=2, high=4, size=3) # (3, 3, 3) box_size = (stack_size - (stack_dim - 1) * margin) / stack_dim for z in range(stack_dim[2]): # Transpose every layer. stack_dim[0], stack_dim[1] = stack_dim[1], stack_dim[0] box_size[0], box_size[1] = box_size[1], box_size[0] for y in range(stack_dim[1]): for x in range(stack_dim[0]): position = (x + 0.5, y + 0.5, z + 0.5) * box_size position[0] += x * margin - stack_size[0] / 2 position[1] += y * margin - stack_size[1] / 2 position[2] += z * margin + 0.03 pose = (position, (0, 0, 0, 1)) pose = utils.multiply(self.zone_pose, pose) urdf = self.fill_template(box_template, {'DIM': box_size}) box_id = env.add_object(urdf, pose) os.remove(urdf) self.color_random_brown(box_id) self.object_points[box_id] = self.get_object_points(box_id) # Randomly select top box on pallet and save ground truth pose. self.goal = {'places': {}, 'steps': []} boxes = env.objects.copy() while boxes: _, height, object_mask = self.get_object_masks(env) top = np.argwhere(height > (np.max(height) - 0.03)) rpixel = top[int(np.floor(np.random.random() * len(top)))] # y, x box_id = int(object_mask[rpixel[0], rpixel[1]]) if box_id in boxes: position, rotation = p.getBasePositionAndOrientation(box_id) rposition = np.float32(position) + np.float32([0, -10, 0]) p.resetBasePositionAndOrientation(box_id, rposition, rotation) self.goal['places'][box_id] = (position, rotation) symmetry = 0 # zone-evaluation: symmetry does not matter self.goal['steps'].append({box_id: (symmetry, [box_id])}) boxes.remove(box_id) self.goal['steps'].reverse() # time-reverse depalletizing self.total_rewards = 0 self.max_steps = len(self.goal['steps']) * 2
def reset(self, env): super().reset(env) # Generate randomly shaped box. box_size = self.get_random_size(0.05, 0.15, 0.05, 0.15, 0.01, 0.06) # Add corner. dimx = (box_size[0] / 2 - 0.025 + 0.0025, box_size[0] / 2 + 0.0025) dimy = (box_size[1] / 2 + 0.0025, box_size[1] / 2 - 0.025 + 0.0025) corner_template = 'assets/corner/corner-template.urdf' replace = {'DIMX': dimx, 'DIMY': dimy} corner_urdf = self.fill_template(corner_template, replace) corner_size = (box_size[0], box_size[1], 0) corner_pose = self.get_random_pose(env, corner_size) env.add_object(corner_urdf, corner_pose, 'fixed') os.remove(corner_urdf) # Add possible placing poses. theta = utils.quatXYZW_to_eulerXYZ(corner_pose[1])[2] fip_rot = utils.eulerXYZ_to_quatXYZW((0, 0, theta + np.pi)) pose1 = (corner_pose[0], fip_rot) alt_x = (box_size[0] / 2) - (box_size[1] / 2) alt_y = (box_size[1] / 2) - (box_size[0] / 2) alt_pos = (alt_x, alt_y, 0) alt_rot0 = utils.eulerXYZ_to_quatXYZW((0, 0, np.pi / 2)) alt_rot1 = utils.eulerXYZ_to_quatXYZW((0, 0, 3 * np.pi / 2)) pose2 = utils.multiply(corner_pose, (alt_pos, alt_rot0)) pose3 = utils.multiply(corner_pose, (alt_pos, alt_rot1)) # Add box. box_template = 'assets/box/box-template.urdf' box_urdf = self.fill_template(box_template, {'DIM': box_size}) box_pose = self.get_random_pose(env, box_size) box_id = env.add_object(box_urdf, box_pose) os.remove(box_urdf) self.color_random_brown(box_id) # Goal: box is aligned with corner (1 of 4 possible poses). self.goals.append(([(box_id, (2 * np.pi, None))], np.int32([[1, 1, 1, 1]]), [corner_pose, pose1, pose2, pose3], False, True, 'pose', None, 1))
def __call__(self, movej, movep, ee, pose0, pose1): """Execute pick and place primitive. Args: movej: function to move robot joints. movep: function to move robot end effector pose. ee: robot end effector. pose0: SE(3) picking pose. pose1: SE(3) placing pose. Returns: timeout: robot movement timed out if True. """ pick_pose, place_pose = pose0, pose1 # Execute picking primitive. prepick_to_pick = ((0, 0, 0.32), (0, 0, 0, 1)) postpick_to_pick = ((0, 0, self.height), (0, 0, 0, 1)) prepick_pose = utils.multiply(pick_pose, prepick_to_pick) postpick_pose = utils.multiply(pick_pose, postpick_to_pick) timeout = movep(prepick_pose) # Move towards pick pose until contact is detected. delta = (np.float32([0, 0, -0.001]), utils.eulerXYZ_to_quatXYZW((0, 0, 0))) targ_pose = prepick_pose while not ee.detect_contact(): # and target_pose[2] > 0: targ_pose = utils.multiply(targ_pose, delta) timeout |= movep(targ_pose) if timeout: return True # Activate end effector, move up, and check picking success. ee.activate() timeout |= movep(postpick_pose, self.speed) pick_success = ee.check_grasp() # Execute placing primitive if pick is successful. if pick_success: preplace_to_place = ((0, 0, self.height), (0, 0, 0, 1)) postplace_to_place = ((0, 0, 0.32), (0, 0, 0, 1)) preplace_pose = utils.multiply(place_pose, preplace_to_place) postplace_pose = utils.multiply(place_pose, postplace_to_place) targ_pose = preplace_pose while not ee.detect_contact(): targ_pose = utils.multiply(targ_pose, delta) timeout |= movep(targ_pose, self.speed) if timeout: return True ee.release() timeout |= movep(postplace_pose) # Move to prepick pose if pick is not successful. else: ee.release() timeout |= movep(prepick_pose) return timeout
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 pick_place(self, pose0, pose1, place_6dof=False): """Execute pick and place primitive. Standard ravens tasks use the `delta` vector to lower the gripper until it makes contact with something. With deformables, however, we need to consider cases when the gripper could detect a rigid OR a soft body (cloth or bag); it should grip the first item it touches. This is handled in the Gripper class. Different deformable ravens tasks use slightly different parameters for better physics (and in some cases, faster simulation). Therefore, rather than make special cases here, those tasks will define their own action parameters, which we use here if they exist. Otherwise, we stick to defaults from standard ravens. Possible action parameters a task might adjust: speed: how fast the gripper moves. delta_z: how fast the gripper lowers for picking / placing. prepick_z: height of the gripper when it goes above the target pose for picking, just before it lowers. postpick_z: after suction gripping, raise to this height, should generally be low for cables / cloth. preplace_z: like prepick_z, but for the placing pose. pause_place: add a small pause for some tasks (e.g., bags) for slightly better soft body physics. final_z: height of the gripper after the action. Recommended to leave it at the default of 0.3, because it has to be set high enough to avoid the gripper occluding the workspace when generating color/depth maps. Args: pose0: picking pose. pose1: placing pose. place_6dof: Boolean to check if we're using 6 DoF placing. Returns: A bool indicating whether the action succeeded or not, via checking the sequence of movep calls. If any movep failed, then self.step() will terminate the episode after this action. """ # Defaults used in the standard Ravens environments. speed = 0.01 delta_z = -0.001 prepick_z = 0.3 postpick_z = 0.3 preplace_z = 0.3 pause_place = 0.0 final_z = 0.3 # Parameters if task provides them, which depends on the task stage. if hasattr(self.task, 'primitive_params'): ts = self.task.task_stage if 'prepick_z' in self.task.primitive_params[ts]: prepick_z = self.task.primitive_params[ts]['prepick_z'] speed = self.task.primitive_params[ts]['speed'] delta_z = self.task.primitive_params[ts]['delta_z'] postpick_z = self.task.primitive_params[ts]['postpick_z'] preplace_z = self.task.primitive_params[ts]['preplace_z'] pause_place = self.task.primitive_params[ts]['pause_place'] # Used to track deformable IDs, so that we can get the vertices. def_ids = [] if self.is_softbody_env(): assert hasattr(self.task, 'def_ids'), 'Soft bodies need def_ids' def_ids = self.task.def_ids # Now proceed as usual with given action parameters. success = True pick_position = np.array(pose0[0]) pick_rotation = np.array(pose0[1]) prepick_position = pick_position.copy() prepick_position[2] = prepick_z # Execute picking motion primitive. prepick_pose = np.hstack((prepick_position, pick_rotation)) success &= self.movep(prepick_pose) target_pose = prepick_pose.copy() delta = np.array([0, 0, delta_z, 0, 0, 0, 0]) # Lower gripper until (a) touch object (rigid OR def), or (a) hit ground. while not self.ee.detect_contact(def_ids) and target_pose[2] > 0: target_pose += delta success &= self.movep(target_pose) # Might need this? if not success: return False # Create constraint (rigid objects) or anchor (deformable). self.ee.activate(self.objects, def_ids) # Increase z slightly (or hard-code it) and check picking success. if self.is_softbody_env() or self.is_new_cable_env(): prepick_pose[2] = postpick_z success &= self.movep(prepick_pose, speed=speed) time.sleep(pause_place) # extra rest for bags elif isinstance(self.task, tasks.names['cable']): prepick_pose[2] = 0.03 success &= self.movep(prepick_pose, speed=0.001) else: prepick_pose[2] += pick_position[2] success &= self.movep(prepick_pose) pick_success = self.ee.check_grasp() if pick_success: place_position = np.array(pose1[0]) place_rotation = np.array(pose1[1]) if not place_6dof: preplace_position = place_position.copy() preplace_position[2] = 0.3 + pick_position[2] preplace_rotation = place_rotation.copy() else: t_world_place = pose1 t_place_preplace = (np.array([0, 0, 0.3]), utils.get_pybullet_quaternion_from_rot( (0, 0, 0))) t_world_preplace = utils.multiply(t_world_place, t_place_preplace) preplace_position = np.array(t_world_preplace[0]) preplace_rotation = np.array(t_world_preplace[1]) # Execute placing motion primitive if pick success. preplace_pose = np.hstack((preplace_position, preplace_rotation)) if self.is_softbody_env() or self.is_new_cable_env(): preplace_pose[2] = preplace_z success &= self.movep(preplace_pose, speed=speed) time.sleep(pause_place) # extra rest for bag elif isinstance(self.task, tasks.names['cable']): preplace_pose[2] = 0.03 success &= self.movep(preplace_pose, speed=0.001) else: success &= self.movep(preplace_pose) # Might need this? # if not success: # return False max_place_steps = 1e3 if not place_6dof: # Lower the gripper. TODO(daniel) test on cloth/bags. target_pose = preplace_pose.copy() place_steps = 0 while not self.ee.detect_contact( def_ids) and target_pose[2] > 0: place_steps += 1 if place_steps > max_place_steps: print('Catching -- I was stuck.') return False target_pose += delta success &= self.movep(target_pose) # Might need this? if not success: return False else: current_pose = preplace_pose.copy() place_pose = np.hstack((place_position, place_rotation)) place_delta = np.linalg.norm(place_pose[0:3] - current_pose[0:3]) place_direction = (place_pose[0:3] - current_pose[0:3]) / place_delta place_steps = 0 while not self.ee.detect_contact( ) and place_delta > PLACE_DELTA_THRESHOLD: place_steps += 1 if place_steps > max_place_steps: print('Catching -- I was stuck.') return False current_pose[0:3] += place_direction * PLACE_STEP success &= self.movep(current_pose) # Might need this? if not success: return False place_delta = np.linalg.norm(place_pose[0:3] - current_pose[0:3]) place_direction = (place_pose[0:3] - current_pose[0:3]) / place_delta # Release and move gripper up (if not 6 DoF) to clear the view for images. self.ee.release() if not place_6dof: preplace_pose[2] = final_z success &= self.movep(preplace_pose) else: # Release and move gripper up to clear the view for images. self.ee.release() prepick_pose[2] = final_z success &= self.movep(prepick_pose) return success
def reset(self, env): super().reset(env) # Add pallet. zone_size = (0.3, 0.25, 0.25) zone_urdf = 'assets/pallet/pallet.urdf' rotation = utils.eulerXYZ_to_quatXYZW((0, 0, 0)) zone_pose = ((0.5, 0.25, 0.02), rotation) env.add_object(zone_urdf, zone_pose, 'fixed') # Add stack of boxes on pallet. margin = 0.01 object_ids = [] object_points = {} stack_size = (0.19, 0.19, 0.19) box_template = 'assets/box/box-template.urdf' stack_dim = np.int32([2, 3, 3]) # stack_dim = np.random.randint(low=2, high=4, size=3) box_size = (stack_size - (stack_dim - 1) * margin) / stack_dim for z in range(stack_dim[2]): # Transpose every layer. stack_dim[0], stack_dim[1] = stack_dim[1], stack_dim[0] box_size[0], box_size[1] = box_size[1], box_size[0] for y in range(stack_dim[1]): for x in range(stack_dim[0]): position = (x + 0.5, y + 0.5, z + 0.5) * box_size position[0] += x * margin - stack_size[0] / 2 position[1] += y * margin - stack_size[1] / 2 position[2] += z * margin + 0.03 pose = (position, (0, 0, 0, 1)) pose = utils.multiply(zone_pose, pose) urdf = self.fill_template(box_template, {'DIM': box_size}) box_id = env.add_object(urdf, pose) os.remove(urdf) object_ids.append((box_id, (0, None))) self.color_random_brown(box_id) object_points[box_id] = self.get_object_points(box_id) # Randomly select top box on pallet and save ground truth pose. targets = [] self.steps = [] boxes = [i[0] for i in object_ids] while boxes: _, height, object_mask = self.get_true_image(env) top = np.argwhere(height > (np.max(height) - 0.03)) rpixel = top[int(np.floor(np.random.random() * len(top)))] # y, x box_id = int(object_mask[rpixel[0], rpixel[1]]) if box_id in boxes: position, rotation = p.getBasePositionAndOrientation(box_id) rposition = np.float32(position) + np.float32([0, -10, 0]) p.resetBasePositionAndOrientation(box_id, rposition, rotation) self.steps.append(box_id) targets.append((position, rotation)) boxes.remove(box_id) self.steps.reverse() # Time-reversed depalletizing. self.goals.append( (object_ids, np.eye(len(object_ids)), targets, False, True, 'zone', (object_points, [(zone_pose, zone_size)]), 1)) self.spawn_box()
def reset(self, env): super().reset(env) # Add container box. zone_size = self.get_random_size(0.05, 0.3, 0.05, 0.3, 0.05, 0.05) zone_pose = self.get_random_pose(env, zone_size) container_template = 'assets/container/container-template.urdf' half = np.float32(zone_size) / 2 replace = {'DIM': zone_size, 'HALF': half} container_urdf = self.fill_template(container_template, replace) env.add_object(container_urdf, zone_pose, 'fixed') os.remove(container_urdf) margin = 0.01 min_object_dim = 0.05 bboxes = [] class TreeNode: def __init__(self, parent, children, bbox): self.parent = parent self.children = children self.bbox = bbox # min x, min y, min z, max x, max y, max z def KDTree(node): size = node.bbox[3:] - node.bbox[:3] # Choose which axis to split. split = size > 2 * min_object_dim if np.sum(split) == 0: bboxes.append(node.bbox) return split = np.float32(split) / np.sum(split) split_axis = np.random.choice(range(len(split)), 1, p=split)[0] # Split along chosen axis and create 2 children cut_ind = np.random.rand() * \ (size[split_axis] - 2 * min_object_dim) + \ node.bbox[split_axis] + min_object_dim child1_bbox = node.bbox.copy() child1_bbox[3 + split_axis] = cut_ind - margin / 2. child2_bbox = node.bbox.copy() child2_bbox[split_axis] = cut_ind + margin / 2. node.children = [ TreeNode(node, [], bbox=child1_bbox), TreeNode(node, [], bbox=child2_bbox) ] KDTree(node.children[0]) KDTree(node.children[1]) # Split container space with KD trees. stack_size = np.array(zone_size) stack_size[0] -= 0.01 stack_size[1] -= 0.01 root_size = (0.01, 0.01, 0) + tuple(stack_size) root = TreeNode(None, [], bbox=np.array(root_size)) KDTree(root) colors = [utils.COLORS[c] for c in utils.COLORS if c != 'brown'] # Add objects in container. object_points = {} object_ids = [] bboxes = np.array(bboxes) object_template = 'assets/box/box-template.urdf' for bbox in bboxes: size = bbox[3:] - bbox[:3] position = size / 2. + bbox[:3] position[0] += -zone_size[0] / 2 position[1] += -zone_size[1] / 2 pose = (position, (0, 0, 0, 1)) pose = utils.multiply(zone_pose, pose) urdf = self.fill_template(object_template, {'DIM': size}) box_id = env.add_object(urdf, pose) os.remove(urdf) object_ids.append((box_id, (0, None))) icolor = np.random.choice(range(len(colors)), 1).squeeze() p.changeVisualShape(box_id, -1, rgbaColor=colors[icolor] + [1]) object_points[box_id] = self.get_object_points(box_id) # Randomly select object in box and save ground truth pose. object_volumes = [] true_poses = [] # self.goal = {'places': {}, 'steps': []} for object_id, _ in object_ids: true_pose = p.getBasePositionAndOrientation(object_id) object_size = p.getVisualShapeData(object_id)[0][3] object_volumes.append(np.prod(np.array(object_size) * 100)) pose = self.get_random_pose(env, object_size) p.resetBasePositionAndOrientation(object_id, pose[0], pose[1]) true_poses.append(true_pose) # self.goal['places'][object_id] = true_pose # symmetry = 0 # zone-evaluation: symmetry does not matter # self.goal['steps'].append({object_id: (symmetry, [object_id])}) # self.total_rewards = 0 # self.max_steps = len(self.goal['steps']) * 2 # Sort oracle picking order by object size. # self.goal['steps'] = [ # self.goal['steps'][i] for i in #. np.argsort(-1 * np.array(object_volumes)) # ] self.goals.append(( object_ids, np.eye(len(object_ids)), true_poses, False, True, 'zone', (object_points, [(zone_pose, zone_size)]), 1))