def get_augmentation_transform(self):
heightmap = np.zeros((320, 160))
theta, trans, pivot = utils.get_random_image_transform_params(
heightmap.shape)
transform_params = theta, trans, pivot
t_world_center, t_world_centeraug = utils.get_se3_from_image_transform(
*transform_params, heightmap, self.bounds, self.pixel_size)
t_worldaug_world = t_world_centeraug @ np.linalg.inv(t_world_center)
return t_worldaug_world, transform_params
def get_six_dof_act(self, transform_params, heightmap, pose0, pose1):
"""Adjust SE(3) poses via the in-plane SE(2) augmentation transform."""
p1_position, p1_rotation = pose1[0], pose1[1]
p0_position, p0_rotation = pose0[0], pose0[1]
if transform_params is not None:
t_world_center, t_world_centernew = utils.get_se3_from_image_transform(
transform_params[0], transform_params[1], transform_params[2],
heightmap, self.bounds, self.pixel_size)
t_worldnew_world = t_world_centernew @ np.linalg.inv(t_world_center)
else:
t_worldnew_world = np.eye(4)
p1_quat_wxyz = (p1_rotation[3], p1_rotation[0], p1_rotation[1],
p1_rotation[2])
t_world_p1 = quaternions.quat2mat(p1_quat_wxyz)
t_world_p1[0:3, 3] = np.array(p1_position)
t_worldnew_p1 = t_worldnew_world @ t_world_p1
p0_quat_wxyz = (p0_rotation[3], p0_rotation[0], p0_rotation[1],
p0_rotation[2])
t_world_p0 = quaternions.quat2mat(p0_quat_wxyz)
t_world_p0[0:3, 3] = np.array(p0_position)
t_worldnew_p0 = t_worldnew_world @ t_world_p0
t_worldnew_p0theta0 = t_worldnew_p0 * 1.0
t_worldnew_p0theta0[0:3, 0:3] = np.eye(3)
# PLACE FRAME, adjusted for this 0 rotation on pick
t_p0_p0theta0 = np.linalg.inv(t_worldnew_p0) @ t_worldnew_p0theta0
t_worldnew_p1theta0 = t_worldnew_p1 @ t_p0_p0theta0
# convert the above rotation to euler
quatwxyz_worldnew_p1theta0 = quaternions.mat2quat(
t_worldnew_p1theta0)
q = quatwxyz_worldnew_p1theta0
quatxyzw_worldnew_p1theta0 = (q[1], q[2], q[3], q[0])
p1_rotation = quatxyzw_worldnew_p1theta0
p1_euler = utils.quatXYZW_to_eulerXYZ(p1_rotation)
roll_scaled = p1_euler[0] / self.theta_scale
pitch_scaled = p1_euler[1] / self.theta_scale
p1_theta_scaled = -p1_euler[2] / self.theta_scale
x = p1_position[0]
y = p1_position[1]
z = p1_position[2]
return np.array([x, y, p1_theta_scaled, roll_scaled, pitch_scaled, z])
def get_data_batch(self, dataset, augment=True):
"""Sample batch."""
batch_obs = []
batch_act = []
for _ in range(self.batch_size):
(obs, act, _, _), _ = dataset.sample()
# Get heightmap from RGB-D images.
configs = self.camera_config
colormap, heightmap = self.get_heightmap(obs, configs)
# self.show_images(colormap, heightmap)
# Concatenate color with depth images.
input_image = np.concatenate(
(colormap, heightmap[Ellipsis, None],
heightmap[Ellipsis, None], heightmap[Ellipsis, None]),
axis=2)
# or just use rgb
# input_image = colormap
# Apply augmentation
if augment:
# note: these pixels are made up,
# just to keep the perturb function happy.
p0 = (160, 80)
p1 = (160, 80)
input_image, _, _, transform_params = utils.perturb(
input_image, [p0, p1], set_theta_zero=False)
t_world_center, t_world_centeraug = utils.get_se3_from_image_transform(
*transform_params, heightmap, self.bounds, self.pixel_size)
t_worldaug_world = t_world_centeraug @ np.linalg.inv(
t_world_center)
else:
t_worldaug_world = np.eye(4)
batch_obs.append(input_image)
batch_act.append(self.act_to_gt_act(
act,
t_worldaug_world)) # this samples pick points from surface
batch_obs = np.array(batch_obs)
batch_act = np.array(batch_act)
return batch_obs, batch_act
def get_six_dof(self,
transform_params,
heightmap,
pose0,
pose1,
augment=True):
"""Adjust SE(3) poses via the in-plane SE(2) augmentation transform."""
debug_visualize = False
p1_position, p1_rotation = pose1[0], pose1[1]
p0_position, p0_rotation = pose0[0], pose0[1]
if debug_visualize:
self.vis = utils.create_visualizer()
self.transport_model.vis = self.vis
if augment:
t_world_center, t_world_centernew = utils.get_se3_from_image_transform(
*transform_params, heightmap, self.bounds, self.pixel_size)
if debug_visualize:
label = 't_world_center'
utils.make_frame(self.vis, label, h=0.05, radius=0.0012, o=1.0)
self.vis[label].set_transform(t_world_center)
label = 't_world_centernew'
utils.make_frame(self.vis, label, h=0.05, radius=0.0012, o=1.0)
self.vis[label].set_transform(t_world_centernew)
t_worldnew_world = t_world_centernew @ np.linalg.inv(t_world_center)
else:
t_worldnew_world = np.eye(4)
p1_quat_wxyz = (p1_rotation[3], p1_rotation[0], p1_rotation[1],
p1_rotation[2])
t_world_p1 = quaternions.quat2mat(p1_quat_wxyz)
t_world_p1[0:3, 3] = np.array(p1_position)
t_worldnew_p1 = t_worldnew_world @ t_world_p1
p0_quat_wxyz = (p0_rotation[3], p0_rotation[0], p0_rotation[1],
p0_rotation[2])
t_world_p0 = quaternions.quat2mat(p0_quat_wxyz)
t_world_p0[0:3, 3] = np.array(p0_position)
t_worldnew_p0 = t_worldnew_world @ t_world_p0
if debug_visualize:
label = 't_worldnew_p1'
utils.make_frame(self.vis, label, h=0.05, radius=0.0012, o=1.0)
self.vis[label].set_transform(t_worldnew_p1)
label = 't_world_p1'
utils.make_frame(self.vis, label, h=0.05, radius=0.0012, o=1.0)
self.vis[label].set_transform(t_world_p1)
label = 't_worldnew_p0-0thetaoriginally'
utils.make_frame(self.vis, label, h=0.05, radius=0.0021, o=1.0)
self.vis[label].set_transform(t_worldnew_p0)
# PICK FRAME, using 0 rotation due to suction rotational symmetry
t_worldnew_p0theta0 = t_worldnew_p0 * 1.0
t_worldnew_p0theta0[0:3, 0:3] = np.eye(3)
if debug_visualize:
label = 'PICK'
utils.make_frame(self.vis, label, h=0.05, radius=0.0021, o=1.0)
self.vis[label].set_transform(t_worldnew_p0theta0)
# PLACE FRAME, adjusted for this 0 rotation on pick
t_p0_p0theta0 = np.linalg.inv(t_worldnew_p0) @ t_worldnew_p0theta0
t_worldnew_p1theta0 = t_worldnew_p1 @ t_p0_p0theta0
if debug_visualize:
label = 'PLACE'
utils.make_frame(self.vis, label, h=0.05, radius=0.0021, o=1.0)
self.vis[label].set_transform(t_worldnew_p1theta0)
# convert the above rotation to euler
quatwxyz_worldnew_p1theta0 = quaternions.mat2quat(t_worldnew_p1theta0)
q = quatwxyz_worldnew_p1theta0
quatxyzw_worldnew_p1theta0 = (q[1], q[2], q[3], q[0])
p1_rotation = quatxyzw_worldnew_p1theta0
p1_euler = utils.quatXYZW_to_eulerXYZ(p1_rotation)
roll = p1_euler[0]
pitch = p1_euler[1]
p1_theta = -p1_euler[2]
p0_theta = 0
z = p1_position[2]
return p0_theta, p1_theta, z, roll, pitch