def _preprocess_poses(self, start_pose, pick_pose, place_pose):
self.start_pose = start_pose
# Slightly lower pick and place z-coords to make contact code register.
mod = ((0, 0, -0.008), (0, 0, 0, 1))
self.pick_pose = utils.multiply(pick_pose, mod)
# self.place_pose = utils.multiply(place_pose, mod)
self.place_pose = place_pose
# Generate intermediate waypoints.
prepick_to_pick = ((0, 0, self.height), (0, 0, 0, 1))
self.prepick_pose = utils.multiply(pick_pose, prepick_to_pick)
postpick_to_pick = ((0, 0, self.height), (0, 0, 0, 1))
self.postpick_pose = utils.multiply(pick_pose, postpick_to_pick)
preplace_to_place = ((0, 0, self.height), (0, 0, 0, 1))
self.preplace_pose = utils.multiply(place_pose, preplace_to_place)
postplace_to_place = ((0, 0, self.height), (0, 0, 0, 1))
self.postplace_pose = utils.multiply(place_pose, postplace_to_place)
self.poses = [
self.start_pose,
self.prepick_pose,
self.pick_pose,
self.postpick_pose,
self.preplace_pose,
self.place_pose,
self.postplace_pose,
]
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 __init__(self):
super().__init__()
self.rot_eps = np.deg2rad(30)
self.train_set = np.int32(
[0, 1, 2, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17, 18, 19])
self.test_set = np.int32([3, 11])
self.homogeneous = True
# Add objects in container.
object_points = {}
object_ids = []
bboxes = np.array(bboxes)
object_template = 'box/box-template-packing.urdf'
self.goal = {'places': {}, 'steps': []}
for bbox in bboxes:
size = bbox[3:] - bbox[:3]
size[2] = size[2] - 0.01
size[1] = size[1] - 0.01
position = size / 2. + bbox[:3]
position[0] += -zone_size[0] / 2
position[1] += -zone_size[1] / 2
position[2] += -zone_size[2] / 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)
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 = '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 = '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:
# SAY: check whether the object arrives its target.
if self.is_match(pose, targs[j], symmetry):
matches[i, :] = 0
matches[:, j] = 0
# Get objects to be picked (prioritize farthest from nearest neighbor).
nn_dists = []
nn_targets = []
for i in range(len(objs)):
object_id, (symmetry, _) = objs[i]
xyz, _ = p.getBasePositionAndOrientation(object_id)
targets_i = np.argwhere(matches[i, :]).reshape(-1)
if len(targets_i) > 0: # pylint: disable=g-explicit-length-test
targets_xyz = np.float32([targs[j][0] for j in targets_i])
dists = np.linalg.norm(
targets_xyz - np.float32(xyz).reshape(1, 3), axis=1)
nn = np.argmin(dists)
nn_dists.append(dists[nn])
nn_targets.append(targets_i[nn])
# Handle ignored objects.
else:
nn_dists.append(0)
nn_targets.append(-1)
order = np.argsort(nn_dists)[::-1]
# SAY: matched objects may be the ones that have been at the target location.
# Filter out matched objects.
order = [i for i in order if nn_dists[i] > 0]
pick_mask = None
for pick_i in order:
pick_mask = np.uint8(obj_mask == objs[pick_i][0])
# Erode to avoid picking on edges.
# pick_mask = cv2.erode(pick_mask, np.ones((3, 3), np.uint8))
if np.sum(pick_mask) > 0:
break
# Trigger task reset if no object is visible.
if pick_mask is None or np.sum(pick_mask) == 0:
self.goals = []
print('Object for pick is not visible. Skipping demonstration.')
return
# Get picking pose.
pick_prob = np.float32(pick_mask)
pick_pix = utils.sample_distribution(pick_prob)
# For "deterministic" demonstrations on insertion-easy, use this:
# pick_pix = (160,80)
pick_pos = utils.pix_to_xyz(pick_pix, hmap,
self.bounds, self.pix_size)
pick_pose = (np.asarray(pick_pos), np.asarray((0, 0, 0, 1)))
# Get placing pose.
targ_pose = targs[nn_targets[pick_i]] # pylint: disable=undefined-loop-variable
obj_pose = p.getBasePositionAndOrientation(objs[pick_i][0]) # pylint: disable=undefined-loop-variable
if not self.sixdof:
obj_euler = utils.quatXYZW_to_eulerXYZ(obj_pose[1])
obj_quat = utils.eulerXYZ_to_quatXYZW((0, 0, obj_euler[2]))
obj_pose = (obj_pose[0], obj_quat)
world_to_pick = utils.invert(pick_pose)
obj_to_pick = utils.multiply(world_to_pick, obj_pose)
pick_to_obj = utils.invert(obj_to_pick)
place_pose = utils.multiply(targ_pose, pick_to_obj)
# Rotate end effector?
if not rotations:
place_pose = (place_pose[0], (0, 0, 0, 1))
place_pose = (np.asarray(place_pose[0]), np.asarray(place_pose[1]))
return {'pose0': pick_pose, 'pose1': place_pose}
def reset(self, env):
super().reset(env)
# Add shelf
# zone_size = self.get_random_size(0.2, 0.45, 0.45, 0.45, 0.1, 0.1)
zone_size = self.get_random_size(0.2, 0.3, 0.2, 0.3, 0.05, 0.05)
zone_pose = self.get_pose_6dof(env, zone_size)
container_template = 'shelf/shelf.urdf'
half = np.float32(zone_size) / 2
rack1 = np.float32(zone_size)
rack2 = (np.float32(zone_size) / 2) - np.float32(zone_size)
replace = {'DIM': zone_size, 'HALF': half, 'RACK': rack1, 'PHR': rack2}
container_urdf = self.fill_template(container_template, replace)
env.add_object(container_urdf, zone_pose, 'fixed')
os.remove(container_urdf)
margin = 0.03
min_object_dim = [0.1, 0.1, 0.1]
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[0] > 2 * min_object_dim[0],
size[1] > 2 * min_object_dim[1],
size[2] > 2 * min_object_dim[2]]
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[split_axis]) + \
node.bbox[split_axis] + min_object_dim[split_axis]
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.
n_objects = 10
obj_shapes = np.random.choice(self.train_set, n_objects)
object_points = {}
object_ids = []
bboxes = np.array(bboxes)
# object_template = 'box/box-template-packing.urdf'
object_template = 'shelf/object-template.urdf'
self.goal = {'places': {}, 'steps': []}
i = 0
for bbox in bboxes:
size = bbox[3:] - bbox[:3]
size[2] = size[2] - 0.01
if size[1] > zone_size[1] / 2.1:
size[1] = (zone_size[1] / 2) - 0.05
position = size / 2. + bbox[:3]
position[0] += -zone_size[0] / 2
position[1] += (-zone_size[1] / 2) + 0.02
position[2] += -zone_size[2] / 2
else:
size[1] = size[1] - 0.01
position = size / 2. + bbox[:3]
position[0] += -zone_size[0] / 2
position[1] += -zone_size[1] / 2
position[2] += -zone_size[2] / 2
pose = (position, (0, 0, 0, 1))
shape = obj_shapes[i]
print ("shape being imported is",obj_shapes[i])
pose = utils.multiply(zone_pose, pose)
fname = f'{shape:02d}.obj'
fname = os.path.join(self.assets_root, 'shelf', fname)
scale = [0.006, 0.002, 0.002]
replace = {'FNAME': (fname,), 'SCALE': scale, 'COLOR': colors[i]}
urdf = self.fill_template(object_template, replace)
block_id = env.add_object(urdf, pose, category='rigid')
os.remove(urdf)
object_ids.append((block_id, (0, None)))
icolor = np.random.choice(range(len(colors)), 1).squeeze()
p.changeVisualShape(block_id, -1, rgbaColor=colors[icolor] + [1])
object_points[block_id] = self.get_object_points(block_id)
i += 1
# 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_shelf(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))
def reset(self, env):
super().reset(env)
# Add pallet.
zone_size = (0.3, 0.25, 0.25)
zone_urdf = '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 = '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 = list((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 = '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 = '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))