def get_target_zone_corners(self): """Determine corners of target zone. Assumes we follow this structure in some order: c2 --- c3 | | | | c1 --- c4 The actual xyz positions depend on the sampled `self.zone_pose`. We use this for (among other things) the target points for the ClothPickPlace demonstrator. Make sure the clockwise ordering is consistent with labels, for the reward and demonstrator. NOTE: the zone stays clockwise in terms of our c1 -> c2 -> c3 -> c4 ordering (from a top down view) but the CLOTH may have been flipped. For now I assume that can be handled in the action by considering the possible counter-clockwise map. """ el2 = self._zone_length / 2 self.c1_position = U.apply(self.zone_pose, (-el2, -el2, 0)) self.c2_position = U.apply(self.zone_pose, (-el2, el2, 0)) self.c3_position = U.apply(self.zone_pose, (el2, el2, 0)) self.c4_position = U.apply(self.zone_pose, (el2, -el2, 0)) self._corner_targets_xy = np.array([ self.c1_position[:2], self.c2_position[:2], self.c3_position[:2], self.c4_position[:2], ])
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.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 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): super().reset(env) # Add stand. base_size = (0.12, 0.36, 0.01) base_urdf = 'assets/hanoi/stand.urdf' base_pose = self.get_random_pose(env, base_size) env.add_object(base_urdf, base_pose, 'fixed') # Rod positions in base coordinates. rod_pos = ((0, -0.12, 0.03), (0, 0, 0.03), (0, 0.12, 0.03)) # Add disks. disks = [] n_disks = 3 for i in range(n_disks): disk_urdf = 'assets/hanoi/disk%d.urdf' % i pos = utils.apply(base_pose, rod_pos[0]) z = 0.015 * (n_disks - i - 2) pos = (pos[0], pos[1], pos[2] + z) disks.append(env.add_object(disk_urdf, (pos, 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(n_disks - 1, 0, 2, 1) # Goal: pick and place disks using Hanoi sequence. for step in hanoi_steps: disk_id = disks[step[0]] targ_pos = rod_pos[step[2]] targ_pos = utils.apply(base_pose, targ_pos) targ_pose = (targ_pos, (0, 0, 0, 1)) self.goals.append( ([(disk_id, (0, None))], np.int32([[1]]), [targ_pose], False, True, 'pose', None, 1 / len(hanoi_steps)))
def reset(self, env): super().reset(env) # Add base. base_size = (0.05, 0.15, 0.005) base_urdf = 'assets/stacking/stand.urdf' base_pose = self.get_random_pose(env, base_size) env.add_object(base_urdf, base_pose, 'fixed') # Block colors. colors = [ utils.COLORS['purple'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow'], utils.COLORS['orange'], utils.COLORS['red'] ] # Add blocks. objs = [] # sym = np.pi / 2 block_size = (0.04, 0.04, 0.04) block_urdf = 'assets/stacking/block.urdf' for i in range(6): block_pose = self.get_random_pose(env, block_size) block_id = env.add_object(block_urdf, block_pose) p.changeVisualShape(block_id, -1, rgbaColor=colors[i] + [1]) objs.append((block_id, (np.pi / 2, None))) # Associate placement locations for goals. place_pos = [(0, -0.05, 0.03), (0, 0, 0.03), (0, 0.05, 0.03), (0, -0.025, 0.08), (0, 0.025, 0.08), (0, 0, 0.13)] targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos] # Goal: blocks are stacked in a pyramid (bottom row: green, blue, purple). self.goals.append((objs[:3], np.ones( (3, 3)), targs[:3], False, True, 'pose', None, 1 / 2)) # Goal: blocks are stacked in a pyramid (middle row: yellow, orange). self.goals.append((objs[3:5], np.ones( (2, 2)), targs[3:5], False, True, 'pose', None, 1 / 3)) # Goal: blocks are stacked in a pyramid (top row: red). self.goals.append((objs[5:], np.ones( (1, 1)), targs[5:], False, True, 'pose', None, 1 / 6))
def reset(self, env): # Add base. base_size = (0.05, 0.15, 0.005) base_urdf = 'assets/stacking/stand.urdf' base_pose = self.random_pose(env, base_size) env.add_object(base_urdf, base_pose, fixed=True) # Block colors. colors = [ utils.COLORS['purple'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow'], utils.COLORS['orange'], utils.COLORS['red'] ] # Add blocks. block_ids = [] block_size = (0.04, 0.04, 0.04) block_urdf = 'assets/stacking/block.urdf' for i in range(6): block_pose = self.random_pose(env, block_size) block_id = env.add_object(block_urdf, block_pose) p.changeVisualShape(block_id, -1, rgbaColor=colors[i] + [1]) block_ids.append(block_id) # Associate placement locations. self.num_steps = 6 self.goal = {'places': {}, 'steps': []} self.goal['places'] = { 0: (utils.apply(base_pose, (0, -0.05, 0.03)), base_pose[1]), 1: (utils.apply(base_pose, (0, 0, 0.03)), base_pose[1]), 2: (utils.apply(base_pose, (0, 0.05, 0.03)), base_pose[1]), 3: (utils.apply(base_pose, (0, -0.025, 0.08)), base_pose[1]), 4: (utils.apply(base_pose, (0, 0.025, 0.08)), base_pose[1]), 5: (utils.apply(base_pose, (0, 0, 0.13)), base_pose[1]) } block_symmetry = np.pi / 2 self.goal['steps'] = [{ block_ids[0]: (block_symmetry, [0, 1, 2]), block_ids[1]: (block_symmetry, [0, 1, 2]), block_ids[2]: (block_symmetry, [0, 1, 2]) }, { block_ids[3]: (block_symmetry, [3, 4]), block_ids[4]: (block_symmetry, [3, 4]) }, { block_ids[5]: (block_symmetry, [5]) }]
def add_cable(self, env, size_range, info): """Add a cable to the env, consisting of rigids beads. Add each bead ID to (a) env.objects, (b) object_points, and (c) cable_bead_ids. Use (b) because the demonstrator checks it to pick the bead farthest from a goal, and it is also used to tally up beads within the zone (to compute reward). Use (c) to distinguish between bead vs non-bead objects in case we add other items. Args: env: A ravens environment. size_range: Used to indicate the area of the target, so the beads avoid spawning there. info: Stores relevant stuff, such as for ground-truth targets. """ num_parts = self.num_parts radius = self.radius length = self.length # Add beaded cable. distance = length / num_parts position, _ = self.random_pose(env, size_range) position = np.float32(position) part_shape = p.createCollisionShape(p.GEOM_BOX, halfExtents=[radius] * 3) part_visual = p.createVisualShape(p.GEOM_SPHERE, radius=radius * 1.5) # Iterate through parts and create constraints as needed. for i in range(num_parts): position[2] += distance parent_frame = (0, 0, distance) part_id = p.createMultiBody(0.1, part_shape, part_visual, basePosition=position) if i > 0: constraint_id = p.createConstraint( parentBodyUniqueId=env.objects[-1], parentLinkIndex=-1, childBodyUniqueId=part_id, childLinkIndex=-1, jointType=p.JOINT_POINT2POINT, jointAxis=(0, 0, 0), parentFramePosition=parent_frame, childFramePosition=(0, 0, 0)) p.changeConstraint(constraint_id, maxForce=100) # Colors if (i > 0) and (i < num_parts - 1): p.changeVisualShape(part_id, -1, rgbaColor=self.color_bead) elif i == num_parts - 1: p.changeVisualShape(part_id, -1, rgbaColor=self.color_end) # Add objects in a consistent manner. self.cable_bead_ids.append(part_id) env.objects.append(part_id) self.object_points[part_id] = np.float32(([0], [0], [0])) # Get target placing positions for each cable bead, if applicable. if (self._name == 'cable-shape' or self._name == 'cable-shape-notarget' or self._name == 'cable-line-notarget'): # ----------------------------------------------------------- # # Here, zone_pose = square_pose, unlike Ravens cable, where the # zone_pose is shifted so that its center matches the straight # line segment center. For `true_position`, we use `zone_pose` # but apply the correct offset to deal with the sides. Note # that `length` is the size of a fully smoothed cable, BUT we # made a rectangle with each side <= length. # ----------------------------------------------------------- # lx = info['lengthx'] ly = info['lengthy'] r = radius if info['nb_sides'] == 1: # Here it's just a straight line on the 'lx' side. x_coord = lx / 2 - (distance * i) y_coord = 0 true_position = (x_coord - r, y_coord, 0) elif info['nb_sides'] == 2: # Start from lx side, go 'left' to the pivot point, then on # the ly side, go 'upwards' but offset by `i`. For radius # offset, I just got this by tuning. XD if i < info['cutoff']: x_coord = lx / 2 - (distance * i) y_coord = -ly / 2 true_position = (x_coord - r, y_coord, 0) else: x_coord = -lx / 2 y_coord = -ly / 2 + (distance * (i - info['cutoff'])) true_position = (x_coord, y_coord + r, 0) elif info['nb_sides'] == 3: # Start from positive lx, positive ly, go down to first # pivot. Then go left to the second pivot, then up again. # For v1, division by two is because we assume BOTH of the # 'ly edges' were divided by two. v1 = (self.num_parts - info['cutoff']) / 2 v2 = self.num_parts - v1 if i < v1: x_coord = lx / 2 y_coord = ly / 2 - (distance * i) true_position = (x_coord, y_coord - r, 0) elif i < v2: x_coord = lx / 2 - (distance * (i - v1)) y_coord = -ly / 2 true_position = (x_coord - r, y_coord, 0) else: x_coord = -lx / 2 y_coord = -ly / 2 + (distance * (i - v2)) true_position = (x_coord, y_coord + r, 0) elif info['nb_sides'] == 4: # I think this is similar to the 2-side case: we start in # the same direction and go counter-clockwise. v1 = info['cutoff'] / 2 v2 = num_parts / 2 v3 = (num_parts + info['cutoff']) / 2 if i < v1: x_coord = lx / 2 - (distance * i) y_coord = -ly / 2 true_position = (x_coord, y_coord, 0) elif i < v2: x_coord = -lx / 2 y_coord = -ly / 2 + (distance * (i - v1)) true_position = (x_coord, y_coord, 0) elif i < v3: x_coord = -lx / 2 + (distance * (i - v2)) y_coord = ly / 2 true_position = (x_coord, y_coord, 0) else: x_coord = lx / 2 y_coord = ly / 2 - (distance * (i - v3)) true_position = (x_coord, y_coord, 0) # Map true_position onto the workspace from zone_pose. true_position = U.apply(self.zone_pose, true_position) # See `cable.py`: just get the places and steps set. self.goal['places'][part_id] = (true_position, (0, 0, 0, 1.)) symmetry = 0 self.goal['steps'][0][part_id] = (symmetry, [part_id]) # Debugging target zones. if self.target_debug_markers: sq_pose = ((true_position[0], true_position[1], 0.002), (0, 0, 0, 1)) sq_template = 'assets/square/square-template-allsides-blue.urdf' replace = {'DIM': (0.003, ), 'HALF': (0.003 / 2, )} urdf = self.fill_template(sq_template, replace) env.add_object(urdf, sq_pose, fixed=True) os.remove(urdf) else: print(f'Warning, env {self._name} will not have goals.')
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)
def reset(self, env): super().reset(env) # Add kit. kit_size = (0.28, 0.2, 0.005) kit_urdf = 'assets/kitting/kit.urdf' kit_pose = self.get_random_pose(env, kit_size) env.add_object(kit_urdf, kit_pose, 'fixed') n_objects = 5 if self.mode == 'train': obj_shapes = np.random.choice(self.train_set, n_objects) else: if self.homogeneous: obj_shapes = [np.random.choice(self.test_set)] * n_objects else: obj_shapes = np.random.choice(self.test_set, n_objects) colors = [ utils.COLORS['purple'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow'], utils.COLORS['red'] ] symmetry = [ 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. targets = [] targ_pos = [[-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]] template = 'assets/kitting/object-template.urdf' for i in range(n_objects): shape = f'{obj_shapes[i]:02d}.obj' scale = [0.003, 0.003, 0.0001] # .0005 pos = utils.apply(kit_pose, targ_pos[i]) theta = np.random.rand() * 2 * np.pi rot = utils.eulerXYZ_to_quatXYZW((0, 0, theta)) replace = {'FNAME': (shape,), 'SCALE': scale, 'COLOR': (0.2, 0.2, 0.2)} urdf = self.fill_template(template, replace) env.add_object(urdf, (pos, rot), 'fixed') os.remove(urdf) targets.append((pos, rot)) # Add objects. objects = [] matches = [] # objects, syms, matcheses = [], [], [] for i in range(n_objects): shape = obj_shapes[i] size = (0.08, 0.08, 0.02) pose = self.get_random_pose(env, size) fname = f'{shape:02d}.obj' scale = [0.003, 0.003, 0.001] # .0005 replace = {'FNAME': (fname,), 'SCALE': scale, 'COLOR': colors[i]} urdf = self.fill_template(template, replace) block_id = env.add_object(urdf, pose) os.remove(urdf) objects.append((block_id, (symmetry[shape], None))) # objects[block_id] = symmetry[shape] match = np.zeros(len(targets)) match[np.argwhere(obj_shapes == shape).reshape(-1)] = 1 matches.append(match) # print(targets) # exit() # matches.append(list(np.argwhere(obj_shapes == shape).reshape(-1))) matches = np.int32(matches) # print(matcheses) # exit() # Add goal. # self.goals.append((objects, syms, targets, 'matches', 'pose', 1.)) # Goal: objects are placed in their respective kit locations. # print(objects) # print(matches) # print(targets) # exit() self.goals.append((objects, matches, targets, False, True, 'pose', None, 1))
def reset(self, env): self.total_rewards = 0 self.goal = {'places': {}, 'steps': [{}]} num_parts = 20 radius = 0.005 length = 2 * radius * num_parts * np.sqrt(2) square_size = (length, length, 0) square_pose = self.random_pose(env, square_size) square_template = 'assets/square/square-template.urdf' replace = {'DIM': (length,), 'HALF': (length / 2 - 0.005,)} urdf = self.fill_template(square_template, replace) env.add_object(urdf, square_pose, fixed=True) os.remove(urdf) # Add goal line. # line_template = 'assets/line/line-template.urdf' self.zone_size = (length, 0.03, 0.2) zone_range = (self.zone_size[0], self.zone_size[1], 0.001) zone_position = (0, length / 2, 0.001) zone_position = utils.apply(square_pose, zone_position) self.zone_pose = (zone_position, square_pose[1]) # urdf = self.fill_template(line_template, {'DIM': (length,)}) # env.add_object(urdf, self.zone_pose, fixed=True) # os.remove(urdf) # Add beaded cable. distance = length / num_parts position, _ = self.random_pose(env, zone_range) position = np.float32(position) part_shape = p.createCollisionShape(p.GEOM_BOX, halfExtents=[radius] * 3) part_visual = p.createVisualShape(p.GEOM_SPHERE, radius=radius * 1.5) # part_visual = p.createVisualShape(p.GEOM_BOX, halfExtents=[radius] * 3) self.object_points = {} for i in range(num_parts): position[2] += distance part_id = p.createMultiBody( 0.1, part_shape, part_visual, basePosition=position) if env.objects: constraint_id = p.createConstraint( parentBodyUniqueId=env.objects[-1], parentLinkIndex=-1, childBodyUniqueId=part_id, childLinkIndex=-1, jointType=p.JOINT_POINT2POINT, jointAxis=(0, 0, 0), parentFramePosition=(0, 0, distance), childFramePosition=(0, 0, 0)) p.changeConstraint(constraint_id, maxForce=100) if (i > 0) and (i < num_parts - 1): color = utils.COLORS['red'] + [1] p.changeVisualShape(part_id, -1, rgbaColor=color) env.objects.append(part_id) self.object_points[part_id] = np.float32((0, 0, 0)).reshape(3, 1) true_position = (radius + distance * i - length / 2, 0, 0) true_position = utils.apply(self.zone_pose, true_position) self.goal['places'][part_id] = (true_position, (0, 0, 0, 1.)) symmetry = 0 # zone-evaluation: symmetry does not matter self.goal['steps'][0][part_id] = (symmetry, [part_id]) # Wait for beaded cable to settle. env.start() time.sleep(1) env.pause()
def reset(self, env): super().reset(env) n_parts = 20 radius = 0.005 length = 2 * radius * n_parts * np.sqrt(2) # Add 3-sided square. square_size = (length, length, 0) square_pose = self.get_random_pose(env, square_size) square_template = 'assets/square/square-template.urdf' replace = {'DIM': (length, ), 'HALF': (length / 2 - 0.005, )} urdf = self.fill_template(square_template, replace) env.add_object(urdf, square_pose, 'fixed') os.remove(urdf) # Get corner points of square. corner0 = (length / 2, length / 2, 0.001) corner1 = (-length / 2, length / 2, 0.001) corner0 = utils.apply(square_pose, corner0) corner1 = utils.apply(square_pose, corner1) # Add cable (series of articulated small blocks). increment = (np.float32(corner1) - np.float32(corner0)) / n_parts position, _ = self.get_random_pose(env, (0.1, 0.1, 0.1)) position = np.float32(position) part_shape = p.createCollisionShape(p.GEOM_BOX, halfExtents=[radius] * 3) part_visual = p.createVisualShape(p.GEOM_SPHERE, radius=radius * 1.5) parent_id = -1 targets = [] objects = [] for i in range(n_parts): position[2] += np.linalg.norm(increment) part_id = p.createMultiBody(0.1, part_shape, part_visual, basePosition=position) if parent_id > -1: constraint_id = p.createConstraint( parentBodyUniqueId=parent_id, parentLinkIndex=-1, childBodyUniqueId=part_id, childLinkIndex=-1, jointType=p.JOINT_POINT2POINT, jointAxis=(0, 0, 0), parentFramePosition=(0, 0, np.linalg.norm(increment)), childFramePosition=(0, 0, 0)) p.changeConstraint(constraint_id, maxForce=100) if (i > 0) and (i < n_parts - 1): color = utils.COLORS['red'] + [1] p.changeVisualShape(part_id, -1, rgbaColor=color) env.obj_ids['rigid'].append(part_id) parent_id = part_id target_xyz = np.float32(corner0) + i * increment + increment / 2 objects.append((part_id, (0, None))) targets.append((target_xyz, (0, 0, 0, 1))) matches = np.clip(np.eye(n_parts) + np.eye(n_parts)[::-1], 0, 1) self.goals.append( (objects, matches, targets, False, False, 'pose', None, 1)) for i in range(480): p.stepSimulation()