def act(self, obs, info): """Run inference and return best action.""" act = {'camera_config': self.camera_config, 'primitive': None} # Get observations and run pick prediction gt_obs = self.info_to_gt_obs(info) pick_prediction = self.pick_model(gt_obs[None, Ellipsis]) if self.use_mdn: pi, mu, var = pick_prediction # prediction = mdn_utils.pick_max_mean(pi, mu, var) pick_prediction = mdn_utils.sample_from_pdf(pi, mu, var) pick_prediction = pick_prediction[:, 0, :] pick_prediction = pick_prediction[0] # unbatch # Get observations and run place prediction obs_with_pick = np.hstack((gt_obs, pick_prediction)) # since the pick at train time is always 0.0, # the predictions are unstable if not exactly 0 obs_with_pick[-1] = 0.0 place_prediction = self.place_model(obs_with_pick[None, Ellipsis]) if self.use_mdn: pi, mu, var = place_prediction # prediction = mdn_utils.pick_max_mean(pi, mu, var) place_prediction = mdn_utils.sample_from_pdf(pi, mu, var) place_prediction = place_prediction[:, 0, :] place_prediction = place_prediction[0] prediction = np.hstack((pick_prediction, place_prediction)) # just go exactly to objects, from observations # p0_position = np.hstack((gt_obs[3:5], 0.02)) # p0_rotation = utils.get_pybullet_quaternion_from_rot( # (0, 0, -gt_obs[5]*self.theta_scale)) # p1_position = np.hstack((gt_obs[0:2], 0.02)) # p1_rotation = utils.get_pybullet_quaternion_from_rot( # (0, 0, -gt_obs[2]*self.theta_scale)) # just go exactly to objects, predicted p0_position = np.hstack((prediction[0:2], 0.02)) p0_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, -prediction[2] * self.theta_scale)) p1_position = np.hstack((prediction[3:5], 0.02)) p1_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, -prediction[5] * self.theta_scale)) # Select task-specific motion primitive. act['primitive'] = 'pick_place' if self.task == 'sweeping': act['primitive'] = 'sweep' elif self.task == 'pushing': act['primitive'] = 'push' params = { 'pose0': (p0_position, p0_rotation), 'pose1': (p1_position, p1_rotation) } act['params'] = params return act
def reset(self, env): self.total_rewards = 0 # Add goal post. line_urdf = 'assets/line/line.urdf' self.zone_size = (0.5, 1, 0.2) self.zone_pose = ((0.5, -0.86, 0), utils.get_pybullet_quaternion_from_rot((0, 0, 0))) env.add_object(line_urdf, self.zone_pose, fixed=True) # Add box. box_size = self.random_size(0.05, 0.25, 0.05, 0.15, 0.05, 0.05) box_template = 'assets/box/box-template.urdf' box_urdf = self.fill_template(box_template, {'DIM': box_size}) px = self.bounds[0, 0] + 0.1 + np.random.rand() * 0.3 position = (px, 0.3, box_size[2] / 2) theta = np.random.rand() * np.pi / 4 - np.pi / 8 rotation = utils.get_pybullet_quaternion_from_rot((0, 0, theta)) box_pose = (position, rotation) box_id = env.add_object(box_urdf, box_pose) os.remove(box_urdf) self.color_random_brown(box_id) self.object_points = {box_id: self.get_object_points(box_id)} # Move end effector to start position next to box. box_dim = p.getVisualShapeData(box_id)[0][3] start_position = np.array([0, box_dim[1] / 2 + 0.01, 0]) start_position = utils.apply(box_pose, start_position) rotation = tuple(env.home_pose[3:]) joints = env.solve_ik((tuple(start_position) + rotation)) for i in range(len(env.joints)): p.resetJointState(env.ur5, env.joints[i], joints[i])
def act(self, obs, gt_act, info): """Run inference and return best action given visual observations.""" del gt_act del info self.regression_model.set_batch_size(1) act = {'camera_config': self.camera_config, 'primitive': None} if not obs: return act # Get heightmap from RGB-D images. colormap, heightmap = self.get_heightmap(obs, self.camera_config) # Concatenate color with depth images. input_image = np.concatenate( (colormap, heightmap[Ellipsis, None], heightmap[Ellipsis, None], heightmap[Ellipsis, None]), axis=2)[None, Ellipsis] # or just use rgb # input_image = colormap[None, ...] # Regression prediction = self.regression_model.forward(input_image) if self.use_mdn: mdn_prediction = prediction pi, mu, var = mdn_prediction # prediction = mdn_utils.pick_max_mean(pi, mu, var) prediction = mdn_utils.sample_from_pdf(pi, mu, var) prediction = prediction[:, 0, :] prediction = prediction[0] p0_position = np.hstack((prediction[0:2], 0.02)) p1_position = np.hstack((prediction[3:5], 0.02)) p0_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, -prediction[2] * self.theta_scale)) p1_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, -prediction[5] * self.theta_scale)) act['primitive'] = 'pick_place' if self.task == 'sweeping': act['primitive'] = 'sweep' elif self.task == 'pushing': act['primitive'] = 'push' params = { 'pose0': (p0_position, p0_rotation), 'pose1': (p1_position, p1_rotation) } act['params'] = params self.regression_model.set_batch_size(self.batch_size) return act
def reset(self, env): # Random base pose rx = env.bounds[0, 0] + 0.2 + np.random.rand() * 0.1 ry = env.bounds[1, 0] + 0.2 + np.random.rand() * 0.6 rtheta = np.random.rand() * 2 * np.pi self.base_pos = np.array([rx, ry, 0.005]) self.base_rot = utils.get_pybullet_quaternion_from_rot((0, 0, rtheta)) base_urdf = 'assets/hanoi/stand.urdf' p.loadURDF(base_urdf, self.base_pos, self.base_rot, useFixedBase=1) # Rod poses in base coordinates self.rod_pos = [[0, -0.12, 0.03], [0, 0, 0.03], [0, 0.12, 0.03]] self.rod_pos = np.array(self.rod_pos) # Add disks env.obj = [] if self.mode == 'train': self.n_disks = np.random.choice([2, 3, 5, 6]) else: self.n_disks = np.random.choice([4, 7]) for i in range(self.n_disks): urdf = 'assets/hanoi/slimdisk.urdf' pos = self.base2origin(self.rod_pos[0, :]) pos[2] += 0.006 * (self.n_disks - i) obj = p.loadURDF(urdf, pos) cw = (self.n_disks - i - 1) / (self.n_disks - 1) color = [cw, 0.6, cw, 1.0] p.changeVisualShape(obj, -1, rgbaColor=color) env.obj.append(obj) self.goal = {'places': {}, 'steps': []} # Solve Hanoi sequence with dynamic programming steps = [] # [[obj id, from rod, to rod], ...] def solve_hanoi(n, t0, t1, t2): if n == 0: steps.append([n, t0, t1]) return solve_hanoi(n - 1, t0, t2, t1) steps.append([n, t0, t1]) solve_hanoi(n - 1, t2, t1, t0) solve_hanoi(self.n_disks - 1, 0, 2, 1) self.n_steps = len(steps) self.max_steps = min(self.n_steps + 10, self.n_steps * 2) # Construct goal sequence [{obj id : (symmetry, pose)}, ...] for step in steps: obj = env.obj[step[0]] pos = self.base2origin(self.rod_pos[step[2], :]) rot = utils.get_pybullet_quaternion_from_rot((0, 0, 0)) place_i = len(self.goal['places']) self.goal['places'][place_i] = (pos, rot) self.goal['steps'].append({obj: (0, [place_i])})
def act(self, obs, info): """Run inference and return best action.""" del obs act = {'camera_config': self.camera_config, 'primitive': None} # Get observations and run predictions. gt_obs = self.info_to_gt_obs(info) # just for visualization gt_act_center = self.info_to_gt_obs(info) # pylint: disable=unused-variable prediction = self.model(gt_obs[None, Ellipsis]) if self.use_mdn: mdn_prediction = prediction pi, mu, var = mdn_prediction # prediction = mdn_utils.pick_max_mean(pi, mu, var) prediction = mdn_utils.sample_from_pdf(pi, mu, var) prediction = prediction[:, 0, :] prediction = prediction[0] # unbatch # self.plot_act_mdn(gt_act_center[None, ...], mdn_prediction) # just go exactly to objects, from observations # p0_position = np.hstack((gt_obs[3:5], 0.02)) # p0_rotation = utils.get_pybullet_quaternion_from_rot( # (0, 0, -gt_obs[5] * self.theta_scale)) # p1_position = np.hstack((gt_obs[0:2], 0.02)) # p1_rotation = utils.get_pybullet_quaternion_from_rot( # (0, 0, -gt_obs[2] * self.theta_scale)) # just go exactly to objects, predicted p0_position = np.hstack((prediction[0:2], 0.02)) p0_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, -prediction[2] * self.theta_scale)) p1_position = np.hstack((prediction[3:5], 0.02)) p1_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, -prediction[5] * self.theta_scale)) # Select task-specific motion primitive. act['primitive'] = 'pick_place' if self.task == 'sweeping': act['primitive'] = 'sweep' elif self.task == 'pushing': act['primitive'] = 'push' params = { 'pose0': (p0_position, p0_rotation), 'pose1': (p1_position, p1_rotation) } act['params'] = params return act
def act(self, obs, info, goal=None): """Run inference and return best action.""" act = {'camera_config': self.camera_config, 'primitive': None} # Get observations and run predictions (second part just for visualization). if self.goal_conditioned: gt_obs = self.info_to_gt_obs(info, goal=goal) gt_act_center = self.info_to_gt_obs(info, goal=goal) else: gt_obs = self.info_to_gt_obs(info) gt_act_center = self.info_to_gt_obs(info) prediction = self.model(gt_obs[None, ...]) if self.USE_MDN: mdn_prediction = prediction pi, mu, var = mdn_prediction #prediction = mdn_utils.pick_max_mean(pi, mu, var) prediction = mdn_utils.sample_from_pdf(pi, mu, var) prediction = prediction[:, 0, :] prediction = prediction[0] # unbatch # Just go exactly to objects, predicted. Daniel: adding 1 rotation inference case. p0_position = np.hstack((prediction[0:2], 0.02)) p0_pred_rot = 0.0 if self.one_rot_inf else -prediction[ 2] * self.THETA_SCALE # idx 2 p0_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, p0_pred_rot)) p1_position = np.hstack((prediction[3:5], 0.02)) p1_pred_rot = 0.0 if self.one_rot_inf else -prediction[ 5] * self.THETA_SCALE # idx 5 p1_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, p1_pred_rot)) # Select task-specific motion primitive. act['primitive'] = 'pick_place' if self.task == 'sweeping': act['primitive'] = 'sweep' elif self.task == 'pushing': act['primitive'] = 'push' params = { 'pose0': (p0_position, p0_rotation), 'pose1': (p1_position, p1_rotation) } act['params'] = params # Daniel: like transporters, determine the task stage if applicable. (AND if loading only) if self.task in ['bag-items-easy', 'bag-items-hard', 'bag-color-goal']: self._determine_task_stage(p0_position, p1_position) return act
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 act(self, obs, info): """Run inference and return best action given visual observations.""" del info act = {'camera_config': self.camera_config, 'primitive': None} if not obs: return act # [Optional] Get heightmap from RGB-D images. colormap, heightmap = self.get_heightmap(obs, self.camera_config) # pylint: disable=unused-variable # Do something here. # Dummy behavior: move to the middle of the workspace. p0_position = (self.bounds[:, 1] - self.bounds[:, 0]) / 2 p0_position += self.bounds[:, 0] p1_position = p0_position rotation = utils.get_pybullet_quaternion_from_rot((0, 0, 0)) # Select task-specific motion primitive. act['primitive'] = 'pick_place' if self.task == 'sweeping': act['primitive'] = 'sweep' elif self.task == 'pushing': act['primitive'] = 'push' params = { 'pose0': (p0_position, rotation), 'pose1': (p1_position, rotation) } act['params'] = params return act
def sweep(self, pose0, pose1): """Execute sweeping primitive.""" success = True position0 = np.float32(pose0[0]) position1 = np.float32(pose1[0]) position0[2] = 0.001 position1[2] = 0.001 direction = position1 - position0 length = np.linalg.norm(position1 - position0) if length == 0: direction = np.float32([0, 0, 0]) else: direction = (position1 - position0) / length theta = np.arctan2(direction[1], direction[0]) rotation = utils.get_pybullet_quaternion_from_rot((0, 0, theta)) over0 = position0.copy() over0[2] = 0.3 over1 = position1.copy() over1[2] = 0.3 success &= self.movep(np.hstack((over0, rotation))) success &= self.movep(np.hstack((position0, rotation))) num_pushes = np.int32(np.floor(length / 0.01)) for _ in range(num_pushes): target = position0 + direction * num_pushes * 0.01 success &= self.movep(np.hstack((target, rotation)), speed=0.003) success &= self.movep(np.hstack((position1, rotation)), speed=0.003) success &= self.movep(np.hstack((over1, rotation))) return success
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 visualize_train_input(self, in_img, p, q, theta, z, roll, pitch): """Visualize the training input.""" points = [] colors = [] height = in_img[:, :, 3] for i in range(in_img.shape[0]): for j in range(in_img.shape[1]): pixel = (i, j) position = utils.pixel_to_position(pixel, height, self.bounds, self.pixel_size) points.append(position) colors.append(in_img[i, j, :3]) points = np.array(points).T # shape (3, N) colors = np.array(colors).T / 255.0 # shape (3, N) self.vis["pointclouds/scene"].set_object( g.PointCloud(position=points, color=colors)) pick_position = utils.pixel_to_position(p, height, self.bounds, self.pixel_size) label = "pick" utils.make_frame(self.vis, label, h=0.05, radius=0.0012, o=0.1) pick_transform = np.eye(4) pick_transform[0:3, 3] = pick_position self.vis[label].set_transform(pick_transform) place_position = utils.pixel_to_position(q, height, self.bounds, self.pixel_size) label = "place" utils.make_frame(self.vis, label, h=0.05, radius=0.0012, o=0.1) place_transform = np.eye(4) place_transform[0:3, 3] = place_position place_transform[2, 3] = z rotation = utils.get_pybullet_quaternion_from_rot( (roll, pitch, -theta)) quaternion_wxyz = np.asarray( [rotation[3], rotation[0], rotation[1], rotation[2]]) place_transform[0:3, 0:3] = mtf.quaternion_matrix(quaternion_wxyz)[0:3, 0:3] self.vis[label].set_transform(place_transform) _, ax = plt.subplots(2, 1) ax[0].imshow(in_img.transpose(1, 0, 2)[:, :, :3] / 255.0) ax[0].scatter(p[0], p[1]) ax[0].scatter(q[0], q[1]) ax[1].imshow(in_img.transpose(1, 0, 2)[:, :, 3]) ax[1].scatter(p[0], p[1]) ax[1].scatter(q[0], q[1]) plt.show()
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 act(self, obs, gt_act, info): """Run inference and return best action.""" del gt_act assert False, 'this needs to have the ordering switched for act inference -- is now xytheta, rpz' # pylint: disable=line-too-long act = {'camera_config': self.camera_config, 'primitive': None} # Get observations and run predictions. gt_obs = self.info_to_gt_obs(info) prediction = self.model(gt_obs[None, Ellipsis]) if self.use_mdn: mdn_prediction = prediction pi, mu, var = mdn_prediction # prediction = mdn_utils.pick_max_mean(pi, mu, var) prediction = mdn_utils.sample_from_pdf(pi, mu, var) prediction = prediction[:, 0, :] prediction = prediction[0] # unbatch p0_position = np.hstack((prediction[0:2], 0.02)) p0_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, -prediction[2] * self.theta_scale)) p1_position = prediction[3:6] p1_rotation = utils.get_pybullet_quaternion_from_rot( (prediction[6] * self.theta_scale, prediction[7] * self.theta_scale, -prediction[8] * self.theta_scale)) # Select task-specific motion primitive. act['primitive'] = 'pick_place_6dof' params = { 'pose0': (p0_position, p0_rotation), 'pose1': (p1_position, p1_rotation) } act['params'] = params return act
def reset(self, env): # Add stand. base_size = (0.12, 0.36, 0.01) base_urdf = 'assets/hanoi/stand.urdf' base_pose = self.random_pose(env, base_size) env.add_object(base_urdf, base_pose, fixed=True) # Rod positions in base coordinates. rod_positions = ((0, -0.12, 0.03), (0, 0, 0.03), (0, 0.12, 0.03)) # Add disks. num_disks = 3 for i in range(num_disks): disk_urdf = 'assets/hanoi/disk%d.urdf' % i position = utils.apply(base_pose, rod_positions[0]) z_offset = 0.015 * (num_disks - i - 2) position = (position[0], position[1], position[2] + z_offset) env.add_object(disk_urdf, (position, base_pose[1])) # Solve Hanoi sequence with dynamic programming. hanoi_steps = [] # [[object index, from rod, to rod], ...] def solve_hanoi(n, t0, t1, t2): if n == 0: hanoi_steps.append([n, t0, t1]) return solve_hanoi(n - 1, t0, t2, t1) hanoi_steps.append([n, t0, t1]) solve_hanoi(n - 1, t2, t1, t0) solve_hanoi(num_disks - 1, 0, 2, 1) self.num_steps = len(hanoi_steps) # Construct goal sequence [{object id : (symmetry, pose)}, ...] self.goal = {'places': {}, 'steps': []} for step in hanoi_steps: object_id = env.objects[step[0]] rod_position = rod_positions[step[2]] place_position = utils.apply(base_pose, rod_position) place_pose = (place_position, utils.get_pybullet_quaternion_from_rot((0, 0, 0))) place_id = len(self.goal['places']) self.goal['places'][place_id] = place_pose self.goal['steps'].append({object_id: (0, [place_id])})
def reset(self, env): self.total_rewards = 0 # Add zone. zone_urdf = 'assets/zone/zone.urdf' self.zone_size = (0.12, 0.12, 0) self.zone_pose = self.random_pose(env, self.zone_size) env.add_object(zone_urdf, self.zone_pose, fixed=True) # Add morsels. self.object_points = {} morsel_urdf = 'assets/morsel/morsel.urdf' for _ in range(50): rx = self.bounds[0, 0] + 0.15 + np.random.rand() * 0.2 ry = self.bounds[1, 0] + 0.4 + np.random.rand() * 0.2 position = (rx, ry, 0.01) theta = np.random.rand() * 2 * np.pi rotation = utils.get_pybullet_quaternion_from_rot((0, 0, theta)) pose = (position, rotation) object_id = env.add_object(morsel_urdf, pose) self.object_points[object_id] = self.get_object_points(object_id)
def act(self, obs, info, goal=None): """Run inference and return best action.""" act = {'camera_config': self.camera_config, 'primitive': None} # Get observations and run pick prediction if self.goal_conditioned: gt_obs = self.info_to_gt_obs(info, goal=goal) else: gt_obs = self.info_to_gt_obs(info) pick_prediction = self.pick_model(gt_obs[None, ...]) if self.USE_MDN: pi, mu, var = pick_prediction #prediction = mdn_utils.pick_max_mean(pi, mu, var) pick_prediction = mdn_utils.sample_from_pdf(pi, mu, var) pick_prediction = pick_prediction[:, 0, :] pick_prediction = pick_prediction[0] # unbatch # Get observations and run place prediction obs_with_pick = np.hstack((gt_obs, pick_prediction)) # since the pick at train time is always 0.0, # the predictions are unstable if not exactly 0 obs_with_pick[-1] = 0.0 place_prediction = self.place_model(obs_with_pick[None, ...]) if self.USE_MDN: pi, mu, var = place_prediction #prediction = mdn_utils.pick_max_mean(pi, mu, var) place_prediction = mdn_utils.sample_from_pdf(pi, mu, var) place_prediction = place_prediction[:, 0, :] place_prediction = place_prediction[0] prediction = np.hstack((pick_prediction, place_prediction)) # Daniel: like with gt_state, guessing this is just for insertion. # just go exactly to objects, from observations # p0_position = np.hstack((gt_obs[3:5], 0.02)) # p0_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -gt_obs[5]*self.THETA_SCALE)) # p1_position = np.hstack((gt_obs[0:2], 0.02)) # p1_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -gt_obs[2]*self.THETA_SCALE)) # Just go exactly to objects, predicted. Daniel: adding 1 rotation inference case. p0_position = np.hstack((prediction[0:2], 0.02)) p0_pred_rot = 0.0 if self.one_rot_inf else -prediction[ 2] * self.THETA_SCALE # idx 2 p0_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, p0_pred_rot)) p1_position = np.hstack((prediction[3:5], 0.02)) p1_pred_rot = 0.0 if self.one_rot_inf else -prediction[ 5] * self.THETA_SCALE # idx 5 p1_rotation = utils.get_pybullet_quaternion_from_rot( (0, 0, p1_pred_rot)) # Select task-specific motion primitive. act['primitive'] = 'pick_place' if self.task == 'sweeping': act['primitive'] = 'sweep' elif self.task == 'pushing': act['primitive'] = 'push' params = { 'pose0': (p0_position, p0_rotation), 'pose1': (p1_position, p1_rotation) } act['params'] = params # Daniel: like transporters, determine the task stage if applicable. (AND if loading only) if self.task in ['bag-items-easy', 'bag-items-hard']: self._determine_task_stage(p0_position, p1_position) return act
def act(self, obs, info, goal=None): """Run inference and return best action given visual observations.""" del info act = {'camera_config': self.camera_config, 'primitive': None} if not obs: return act # Get heightmap from RGB-D images. colormap, heightmap = self.get_heightmap(obs, self.camera_config) if goal is not None: colormap_g, heightmap_g = self.get_heightmap( goal, self.camera_config) # Concatenate color with depth images. input_image = self.concatenate_c_h(colormap, heightmap) # Make a goal image if needed, and for consistency stack with input. if self.use_goal_image: goal_image = self.concatenate_c_h(colormap_g, heightmap_g) input_image = np.concatenate((input_image, goal_image), axis=2) assert input_image.shape[2] == 12, input_image.shape # Attention model forward pass. if self.use_goal_image: half = int(input_image.shape[2] / 2) input_only = input_image[:, :, :half] # ignore goal portion attention = self.attention_model.forward(input_only) else: attention = self.attention_model.forward(input_image) argmax = np.argmax(attention) argmax = np.unravel_index(argmax, shape=attention.shape) p0_pixel = argmax[:2] p0_theta = argmax[2] * (2 * np.pi / attention.shape[2]) # Transport model forward pass. if isinstance(self.transport_model, TransportGoal): half = int(input_image.shape[2] / 2) img_curr = input_image[:, :, :half] img_goal = input_image[:, :, half:] transport = self.transport_model.forward(img_curr, img_goal, p0_pixel) else: transport = self.transport_model.forward(input_image, p0_pixel) argmax = np.argmax(transport) argmax = np.unravel_index(argmax, shape=transport.shape) p1_pixel = argmax[:2] p1_theta = argmax[2] * (2 * np.pi / transport.shape[2]) # Pixels to end effector poses. p0_position = utils.pixel_to_position(p0_pixel, heightmap, self.bounds, self.pixel_size) p1_position = utils.pixel_to_position(p1_pixel, heightmap, self.bounds, self.pixel_size) p0_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -p0_theta)) p1_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -p1_theta)) return self.p0_p1_position_rotations_to_act(act, p0_position, p0_rotation, p1_position, p1_rotation)
def act(self, obs, info, compute_error=False, gt_act=None): """Run inference and return best action given visual observations.""" act = {'camera_config': self.camera_config, 'primitive': None} if not obs: return act # Get heightmap from RGB-D images. colormap, heightmap = self.get_heightmap(obs, self.camera_config) # Concatenate color with depth images. input_image = np.concatenate( (colormap, heightmap[Ellipsis, None], heightmap[Ellipsis, None], heightmap[Ellipsis, None]), axis=2) # Attention model forward pass. attention = self.attention_model.forward(input_image) argmax = np.argmax(attention) argmax = np.unravel_index(argmax, shape=attention.shape) p0_pixel = argmax[:2] p0_theta = argmax[2] * (2 * np.pi / attention.shape[2]) # Transport model forward pass. transport = self.transport_model.forward(input_image, p0_pixel) _, z, roll, pitch = self.rpz_model.forward(input_image, p0_pixel) argmax = np.argmax(transport) argmax = np.unravel_index(argmax, shape=transport.shape) # Index into 3D discrete tensor, grab z, roll, pitch activations z_best = z[:, argmax[0], argmax[1], argmax[2]][Ellipsis, None] roll_best = roll[:, argmax[0], argmax[1], argmax[2]][Ellipsis, None] pitch_best = pitch[:, argmax[0], argmax[1], argmax[2]][Ellipsis, None] # Send through regressors for each of z, roll, pitch z_best = self.rpz_model.z_regressor(z_best)[0, 0] roll_best = self.rpz_model.roll_regressor(roll_best)[0, 0] pitch_best = self.rpz_model.pitch_regressor(pitch_best)[0, 0] p1_pixel = argmax[:2] p1_theta = argmax[2] * (2 * np.pi / transport.shape[2]) # Pixels to end effector poses. p0_position = utils.pixel_to_position(p0_pixel, heightmap, self.bounds, self.pixel_size) p1_position = utils.pixel_to_position(p1_pixel, heightmap, self.bounds, self.pixel_size) p1_position = (p1_position[0], p1_position[1], z_best) p0_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -p0_theta)) p1_rotation = utils.get_pybullet_quaternion_from_rot( (roll_best, pitch_best, -p1_theta)) if compute_error: gt_p0_position, gt_p0_rotation = gt_act['params']['pose0'] gt_p1_position, gt_p1_rotation = gt_act['params']['pose1'] gt_p0_pixel = np.array( utils.position_to_pixel(gt_p0_position, self.bounds, self.pixel_size)) gt_p1_pixel = np.array( utils.position_to_pixel(gt_p1_position, self.bounds, self.pixel_size)) self.p0_pixel_error( np.linalg.norm(gt_p0_pixel - np.array(p0_pixel))) self.p1_pixel_error( np.linalg.norm(gt_p1_pixel - np.array(p1_pixel))) gt_p0_theta = -np.float32( utils.get_rot_from_pybullet_quaternion(gt_p0_rotation)[2]) gt_p1_theta = -np.float32( utils.get_rot_from_pybullet_quaternion(gt_p1_rotation)[2]) self.p0_theta_error( abs((np.rad2deg(gt_p0_theta - p0_theta) + 180) % 360 - 180)) self.p1_theta_error( abs((np.rad2deg(gt_p1_theta - p1_theta) + 180) % 360 - 180)) return None return self.p0_p1_position_rotations_to_act(act, p0_position, p0_rotation, p1_position, p1_rotation)
def random_pose(self, env, object_size): position, rotation = super(InsertionEasy, self).random_pose(env, object_size) rotation = utils.get_pybullet_quaternion_from_rot((0, 0, np.pi / 2)) return position, rotation
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 act(self, obs, info): """Run inference and return best action.""" act = {'camera_config': self.camera_config, 'primitive': None} # Get observations and run pick prediction gt_obs = self.info_to_gt_obs(info) pick_prediction = self.pick_model(gt_obs[None, Ellipsis]) if self.use_mdn: pi, mu, var = pick_prediction # prediction = mdn_utils.pick_max_mean(pi, mu, var) pick_prediction = mdn_utils.sample_from_pdf(pi, mu, var) pick_prediction = pick_prediction[:, 0, :] pick_prediction = pick_prediction[0] # unbatch # Get observations and run place prediction obs_with_pick = np.hstack((gt_obs, pick_prediction)).astype(np.float32) # since the pick at train time is always 0.0, # the predictions are unstable if not exactly 0 obs_with_pick[-1] = 0.0 place_se2_prediction = self.place_se2_model(obs_with_pick[None, Ellipsis]) if self.use_mdn: pi, mu, var = place_se2_prediction # prediction = mdn_utils.pick_max_mean(pi, mu, var) place_se2_prediction = mdn_utils.sample_from_pdf(pi, mu, var) place_se2_prediction = place_se2_prediction[:, 0, :] place_se2_prediction = place_se2_prediction[0] # Get observations and run rpz prediction obs_with_pick_place_se2 = np.hstack( (obs_with_pick, place_se2_prediction)).astype(np.float32) place_rpz_prediction = self.place_rpz_model( obs_with_pick_place_se2[None, Ellipsis]) if self.use_mdn: pi, mu, var = place_rpz_prediction # prediction = mdn_utils.pick_max_mean(pi, mu, var) place_rpz_prediction = mdn_utils.sample_from_pdf(pi, mu, var) place_rpz_prediction = place_rpz_prediction[:, 0, :] place_rpz_prediction = place_rpz_prediction[0] p0_position = np.hstack((pick_prediction[0:2], 0.02)) p0_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, 0)) p1_position = np.hstack( (place_se2_prediction[0:2], place_rpz_prediction[2])) p1_rotation = utils.get_pybullet_quaternion_from_rot( (place_rpz_prediction[0] * self.theta_scale, place_rpz_prediction[1] * self.theta_scale, -place_se2_prediction[2] * self.theta_scale)) # Select task-specific motion primitive. act['primitive'] = 'pick_place_6dof' params = { 'pose0': (p0_position, p0_rotation), 'pose1': (p1_position, p1_rotation) } act['params'] = params return act
def act(self, obs, info): """Run inference and return best action given visual observations.""" del info act = {'camera_config': self.camera_config, 'primitive': None} if not obs: return act # Get heightmap from RGB-D images. colormap, heightmap = self.get_heightmap(obs, self.camera_config) # Concatenate color with depth images. input_image = np.concatenate( (colormap, heightmap[Ellipsis, None], heightmap[Ellipsis, None], heightmap[Ellipsis, None]), axis=2) # Get top-k pixels from pick and place heatmaps. k = 100 pick_heatmap = self.pick_model.forward(input_image, apply_softmax=True).squeeze() place_heatmap = self.place_model.forward(input_image, apply_softmax=True).squeeze() descriptors = np.float32(self.match_model.forward(input_image)) # V4 pick_heatmap = cv2.GaussianBlur(pick_heatmap, (49, 49), 0) place_heatmap = cv2.GaussianBlur(place_heatmap, (49, 49), 0) pick_topk = np.int32( np.unravel_index( np.argsort(pick_heatmap.reshape(-1))[-k:], pick_heatmap.shape)).T pick_pixel = pick_topk[-1, :] from skimage.feature import peak_local_max # pylint: disable=g-import-not-at-top place_peaks = peak_local_max(place_heatmap, num_peaks=1) distances = np.ones((place_peaks.shape[0], self.num_rotations)) * 10 pick_descriptor = descriptors[0, pick_pixel[0], pick_pixel[1], :].reshape(1, -1) for i in range(place_peaks.shape[0]): peak = place_peaks[i, :] place_descriptors = descriptors[:, peak[0], peak[1], :] distances[i, :] = np.linalg.norm(place_descriptors - pick_descriptor, axis=1) ibest = np.unravel_index(np.argmin(distances), shape=distances.shape) p0_pixel = pick_pixel p0_theta = 0 p1_pixel = place_peaks[ibest[0], :] p1_theta = ibest[1] * (2 * np.pi / self.num_rotations) # # V3 # pick_heatmap = cv2.GaussianBlur(pick_heatmap, (49, 49), 0) # place_heatmap = cv2.GaussianBlur(place_heatmap, (49, 49), 0) # pick_topk = np.int32( # np.unravel_index( # np.argsort(pick_heatmap.reshape(-1))[-k:], pick_heatmap.shape)).T # place_topk = np.int32( # np.unravel_index( # np.argsort(place_heatmap.reshape(-1))[-k:], # place_heatmap.shape)).T # pick_pixel = pick_topk[-1, :] # place_pixel = place_topk[-1, :] # pick_descriptor = descriptors[0, pick_pixel[0], # pick_pixel[1], :].reshape(1, -1) # place_descriptor = descriptors[:, place_pixel[0], place_pixel[1], :] # distances = np.linalg.norm(place_descriptor - pick_descriptor, axis=1) # irotation = np.argmin(distances) # p0_pixel = pick_pixel # p0_theta = 0 # p1_pixel = place_pixel # p1_theta = irotation * (2 * np.pi / self.num_rotations) # # V2 # pick_topk = np.int32( # np.unravel_index( # np.argsort(pick_heatmap.reshape(-1))[-k:], pick_heatmap.shape)).T # place_topk = np.int32( # np.unravel_index( # np.argsort(place_heatmap.reshape(-1))[-k:], # place_heatmap.shape)).T # pick_pixel = pick_topk[-1, :] # pick_descriptor = descriptors[0, pick_pixel[0], # pick_pixel[1], :].reshape(1, 1, 1, -1) # distances = np.linalg.norm(descriptors - pick_descriptor, axis=3) # distances = np.transpose(distances, [1, 2, 0]) # max_distance = int(np.round(np.max(distances))) # for i in range(self.num_rotations): # distances[:, :, i] = cv2.circle(distances[:, :, i], # (pick_pixel[1], pick_pixel[0]), 50, # max_distance, -1) # ibest = np.unravel_index(np.argmin(distances), shape=distances.shape) # p0_pixel = pick_pixel # p0_theta = 0 # p1_pixel = ibest[:2] # p1_theta = ibest[2] * (2 * np.pi / self.num_rotations) # # V1 # pick_topk = np.int32( # np.unravel_index( # np.argsort(pick_heatmap.reshape(-1))[-k:], pick_heatmap.shape)).T # place_topk = np.int32( # np.unravel_index( # np.argsort(place_heatmap.reshape(-1))[-k:], # place_heatmap.shape)).T # distances = np.zeros((k, k, self.num_rotations)) # for ipick in range(k): # pick_descriptor = descriptors[0, pick_topk[ipick, 0], # pick_topk[ipick, 1], :].reshape(1, -1) # for iplace in range(k): # place_descriptors = descriptors[:, place_topk[iplace, 0], # place_topk[iplace, 1], :] # distances[ipick, iplace, :] = np.linalg.norm( # place_descriptors - pick_descriptor, axis=1) # ibest = np.unravel_index(np.argmin(distances), shape=distances.shape) # p0_pixel = pick_topk[ibest[0], :] # p0_theta = 0 # p1_pixel = place_topk[ibest[1], :] # p1_theta = ibest[2] * (2 * np.pi / self.num_rotations) # Pixels to end effector poses. p0_position = utils.pixel_to_position(p0_pixel, heightmap, self.bounds, self.pixel_size) p1_position = utils.pixel_to_position(p1_pixel, heightmap, self.bounds, self.pixel_size) p0_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -p0_theta)) p1_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -p1_theta)) act['primitive'] = 'pick_place' if self.task == 'sweeping': act['primitive'] = 'sweep' elif self.task == 'pushing': act['primitive'] = 'push' params = { 'pose0': (p0_position, p0_rotation), 'pose1': (p1_position, p1_rotation) } act['params'] = params return act
def reset(self, env): # Add kit. kit_size = (0.28, 0.2, 0.005) kit_urdf = 'assets/kitting/kit.urdf' kit_pose = self.random_pose(env, kit_size) env.add_object(kit_urdf, kit_pose, fixed=True) # num_shapes = 20 # train_split = 14 num_objects = 5 if self.mode == 'train': object_shapes = np.random.choice(self.train_set, num_objects) else: if self.homogeneous: object_shapes = [np.random.choice(self.test_set)] * num_objects else: object_shapes = np.random.choice(self.test_set, num_objects) self.num_steps = num_objects self.goal = {'places': {}, 'steps': [{}]} colors = [ utils.COLORS['purple'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow'], utils.COLORS['red'] ] symmetries = [ 2 * np.pi, 2 * np.pi, 2 * np.pi / 3, np.pi / 2, np.pi / 2, 2 * np.pi, np.pi, 2 * np.pi / 5, np.pi, np.pi / 2, 2 * np.pi / 5, 0, 2 * np.pi, 2 * np.pi, 2 * np.pi, 2 * np.pi, 0, 2 * np.pi / 6, 2 * np.pi, 2 * np.pi ] # Build kit. place_positions = [[-0.09, 0.045, 0.0014], [0, 0.045, 0.0014], [0.09, 0.045, 0.0014], [-0.045, -0.045, 0.0014], [0.045, -0.045, 0.0014]] object_template = 'assets/kitting/object-template.urdf' for i in range(num_objects): shape = f'{object_shapes[i]:02d}.obj' scale = [0.003, 0.003, 0.0001] # .0005 position = utils.apply(kit_pose, place_positions[i]) theta = np.random.rand() * 2 * np.pi rotation = utils.get_pybullet_quaternion_from_rot((0, 0, theta)) replace = { 'FNAME': (shape, ), 'SCALE': scale, 'COLOR': (0.2, 0.2, 0.2) } urdf = self.fill_template(object_template, replace) _ = env.add_object(urdf, (position, rotation), fixed=True) os.remove(urdf) self.goal['places'][i] = (position, rotation) # Add objects. for i in range(num_objects): ishape = object_shapes[i] size = (0.08, 0.08, 0.02) pose = self.random_pose(env, size) shape = f'{ishape:02d}.obj' scale = [0.003, 0.003, 0.001] # .0005 replace = {'FNAME': (shape, ), 'SCALE': scale, 'COLOR': colors[i]} urdf = self.fill_template(object_template, replace) block_id = env.add_object(urdf, pose) os.remove(urdf) places = np.argwhere(ishape == object_shapes).reshape(-1) self.goal['steps'][0][block_id] = (symmetries[ishape], places)