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)
