def __init__(self, name, task):
# self.fig, self.ax = plt.subplots(1,1)
self.name = name
self.task = task
if self.task in ['aligning', 'palletizing', 'packing']:
self.use_box_dimensions = True
else:
self.use_box_dimensions = False
if self.task in ['sorting']:
self.use_colors = True
else:
self.use_colors = False
self.total_iter = 0
self.camera_config = cameras.RealSenseD415.CONFIG
# A place to save pre-trained models.
self.models_dir = os.path.join('checkpoints', self.name)
if not tf.io.gfile.exists(self.models_dir):
os.makedirs(self.models_dir)
# Set up model.
self.model = None
# boundaries = [1000, 2000, 5000, 10000]
# values = [1e-2, 1e-2, 1e-2, 1e-3, 1e-5]
# learning_rate_fn = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
# boundaries, values)
# self.optim = tf.keras.optimizers.Adam(learning_rate=learning_rate_fn)
self.optim = tf.keras.optimizers.Adam(learning_rate=2e-4)
self.metric = tf.keras.metrics.Mean(name='metric')
self.val_metric = tf.keras.metrics.Mean(name='val_metric')
self.theta_scale = 10.0
self.batch_size = 128
self.use_mdn = True
self.vis = utils.create_visualizer()
self.bounds = np.array([[0.25, 0.75], [-0.5, 0.5], [0, 0.28]])
self.pixel_size = 0.003125
self.six_dof = False
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