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