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))
