def decode_base_angle_encoding(self, pick_samples):
base_angles = pick_samples[:, 4:6]
base_angles = [
utils.decode_sin_and_cos_to_angle(base_angle)
for base_angle in base_angles
]
pick_samples[:, 4] = base_angles
pick_samples = np.delete(
pick_samples, 5, 1) # remove unnecessary 2-dim encoded base angle
return pick_samples
def get_pick_base_poses(action, smples):
pick_base_poses = []
for smpl in smples:
smpl = smpl[0:4]
sin_cos_encoding = smpl[-2:]
decoded_angle = utils.decode_sin_and_cos_to_angle(sin_cos_encoding)
smpl = np.hstack([smpl[0:2], decoded_angle])
abs_base_pose = utils.get_absolute_pick_base_pose_from_ir_parameters(
smpl, action.discrete_parameters['object'])
pick_base_poses.append(abs_base_pose)
return pick_base_poses
def get_place_base_poses(action, smples, mover):
place_base_poses = smples[:, 4:]
to_return = []
for bsmpl in place_base_poses:
sin_cos_encoding = bsmpl[-2:]
decoded_angle = utils.decode_sin_and_cos_to_angle(sin_cos_encoding)
bsmpl = np.hstack([bsmpl[0:2], decoded_angle])
to_return.append(bsmpl)
to_return = np.array(to_return)
to_return[:, 0:2] += mover.regions[
action.discrete_parameters['region']].box[0]
return to_return
def unprocess_pick_smpls(smpls):
unprocessed = []
for smpl in smpls:
grasp_params = smpl[0:3]
ir_parameters = smpl[3:]
portion = ir_parameters[0]
base_angle = utils.decode_sin_and_cos_to_angle(ir_parameters[1:3])
facing_angle_offset = ir_parameters[3]
unprocessed.append(
np.hstack([grasp_params, portion, base_angle,
facing_angle_offset]))
return np.array(unprocessed)
def unprocess_pick_and_place_smpls(smpls):
pick_smpls = smpls[0]
place_smpls = smpls[1]
pick_unprocessed = []
place_unprocessed = []
for pick_smpl, place_smpl in zip(pick_smpls, place_smpls):
grasp_params = pick_smpl[0:3]
ir_parameters = pick_smpl[3:]
portion = ir_parameters[0]
base_angle = utils.decode_sin_and_cos_to_angle(ir_parameters[1:3])
facing_angle_offset = ir_parameters[3]
pick_unprocessed.append(
np.hstack([grasp_params, portion, base_angle,
facing_angle_offset]))
abs_base_pose = utils.decode_pose_with_sin_and_cos_angle(place_smpl)
place_unprocessed.append(abs_base_pose)
smpls = np.hstack([pick_unprocessed, place_unprocessed])
return smpls
def generate_pick_and_place_batch(smpler_state, policy, noise_batch):
pick_smpler = policy['pick']
inp = prepare_input(smpler_state,
noise_batch,
delete=True,
region='loading_region',
filter_konfs=False)
pick_samples = pick_smpler.policy_model.predict(inp)
pick_base_poses = []
for p in pick_samples:
ir_parameters = p[3:]
portion = ir_parameters[0]
base_angle = utils.decode_sin_and_cos_to_angle(ir_parameters[1:3])
facing_angle_offset = ir_parameters[3]
pick_param = np.hstack(
[p[:3], portion, base_angle, facing_angle_offset])
_, pick_base_pose = utils.get_pick_base_pose_and_grasp_from_pick_parameters(
smpler_state.obj, pick_param)
pick_base_pose = utils.encode_pose_with_sin_and_cos_angle(
pick_base_pose)
pick_base_poses.append(pick_base_pose)
pick_base_poses = np.array(pick_base_poses)
inp = prepare_input(smpler_state,
noise_batch,
delete=True,
region=smpler_state.region,
filter_konfs=False)
poses = inp[-2]
poses[:, -4:] = pick_base_poses
inp[-2] = poses
z_smpls = gaussian_noise(z_size=(len(noise_batch), 4))
inp[-1] = z_smpls
place_smpler = policy['place']
place_samples = place_smpler.policy_model.predict(inp)
return pick_samples, place_samples