def sample_object_position(T_aug=None,
upright=False):
pos = np.array([0, 0, 0.1])
quat = None
if upright:
quat = np.array([1,0,0,0])
else:
quat = transforms3d.euler.euler2quat(np.deg2rad(90), 0, 0)
T_O_slider = transform_utils.transform_from_pose(pos, quat)
# apply a random yaw to the object
yaw = OBJECT_YAW
quat_yaw = transforms3d.euler.euler2quat(0, 0, yaw)
T_yaw = transform_utils.transform_from_pose([0,0,0], quat_yaw)
T_O_slider = np.matmul(T_yaw, T_O_slider)
T_W_slider = None
if T_aug is not None:
T_W_slider = T_aug @ T_O_slider
else:
T_W_slider = T_O_slider
pose_dict = transform_utils.matrix_to_dict(T_W_slider)
# note the quat/pos ordering
q = np.concatenate((pose_dict['quaternion'], pose_dict['position']))
return q
def sample_slider_position(T_aug=None):
# nominal position with box lying on it's side.
# the side with the recipe for "Colored Sugar" is facing upwards
pos = np.array([1.56907481e-04, 1.11390697e-06, 5.11972761e-02])
quat = np.array(
[7.13518047e-01, -6.69765583e-07, -7.00636851e-01, -6.82079212e-07])
T_O_slider = transform_utils.transform_from_pose(pos, quat)
# apply a random yaw to the object
yaw_min = 0
yaw_max = 2 * np.pi
yaw = random_sample_in_range(0, 2 * np.pi)
quat_yaw = transforms3d.euler.euler2quat(0, 0, yaw)
T_yaw = transform_utils.transform_from_pose([0, 0, 0], quat_yaw)
T_O_slider = np.matmul(T_yaw, T_O_slider)
T_W_slider = None
if T_aug is not None:
T_W_slider = T_aug @ T_O_slider
else:
T_W_slider = T_O_slider
pose_dict = transform_utils.matrix_to_dict(T_W_slider)
# note the quat/pos ordering
q = np.concatenate((pose_dict['quaternion'], pose_dict['position']))
return q
def get_initial_state():
# set initial condition
slider_pos = np.array([0.0, 0.0, 0.3])
slider_quat = np.array([1, 0, 0, 0])
if BOX_ON_SIDE:
slide_pos = np.array([1.56907481e-04, 1.11390697e-06, 5.11972761e-02])
slider_quat = np.array([
7.13518047e-01, -6.69765583e-07, -7.00636851e-01, -6.82079212e-07
])
T_O_slider = transform_utils.transform_from_pose(slide_pos, slider_quat)
quat_yaw = transforms3d.euler.euler2quat(0, 0, OBJECT_YAW)
T_yaw = transform_utils.transform_from_pose([0, 0, 0], quat_yaw)
T_O_slider = np.matmul(T_yaw, T_O_slider)
# O denotes the canonical frame
pusher_pos_homog_O = np.array([-0.15, 0, 0, 1])
# pusher_pos_homog_O = np.array([-0.075, 0, 0, 1])
pusher_velocity_3d_O = PUSHER_VELOCITY * np.array(
[np.cos(PUSHER_ANGLE), np.sin(PUSHER_ANGLE), 0])
# apply transform T_W_O to pusher and slider
T_W_O = np.eye(4)
T_W_O[:3, 3] = GLOBAL_TRANSLATION
T_W_O[:3, :3] = transforms3d.euler.euler2mat(0, 0, YAW)
T_W_slider = T_W_O @ T_O_slider
pusher_pos_3d_W = T_W_O @ pusher_pos_homog_O
pusher_velocity_3d_W = (T_W_O[:3, :3] @ pusher_velocity_3d_O)
# slider
slider_pose_dict = transform_utils.matrix_to_dict(T_W_slider)
q_slider = np.concatenate(
[slider_pose_dict['quaternion'], slider_pose_dict['position']])
q_pusher = pusher_pos_3d_W[:2]
action = pusher_velocity_3d_W[:2]
# expand it to be an action sequence
# [N, 2]
action_sequence = torch.Tensor(action).unsqueeze(0).expand([N, -1])
return {
'T_W_O': T_W_O,
'pusher_pos_3d': pusher_pos_3d_W,
'T_W_slider': T_W_slider,
'q_slider': q_slider,
'q_pusher': q_pusher,
'action': action,
'action_sequence': action_sequence,
}
def sample_slider_position(T_aug=None):
# always at the origin basically
pos = np.array([0, 0, 0.03])
yaw_min = np.array([0])
yaw_max = np.array([2 * np.pi])
yaw_sampler = Box(yaw_min, yaw_max)
yaw = yaw_sampler.sample()
quat = transforms3d.euler.euler2quat(0, 0, yaw)
T_O_slider = transform_utils.transform_from_pose(pos, quat)
T_W_slider = None
if T_aug is not None:
T_W_slider = T_aug @ T_O_slider
else:
T_W_slider = T_O_slider
pose_dict = transform_utils.matrix_to_dict(T_W_slider)
# note the quat/pos orderining
q = np.concatenate((pose_dict['quaternion'], pose_dict['position']))
return q
def get_initial_state():
# set initial condition
slider_pos = np.array([0.0, 0.0, 0.3])
slider_quat = np.array([1, 0, 0, 0])
slider_quat = transforms3d.euler.euler2quat(0, 0, OBJECT_YAW)
# O denotes the canonical frame
T_O_slider = transform_utils.transform_from_pose(slider_pos, slider_quat)
# pusher_pos_homog_O = np.array([-0.15, 0, 0, 1])
pusher_pos_homog_O = np.array([-0.075, 0, 0, 1])
pusher_velocity_3d_O = PUSHER_VELOCITY * np.array(
[np.cos(PUSHER_ANGLE), np.sin(PUSHER_ANGLE), 0])
# apply transform T_W_O to pusher and slider
T_W_O = np.eye(4)
T_W_O[:3, 3] = GLOBAL_TRANSLATION
T_W_O[:3, :3] = transforms3d.euler.euler2mat(0, 0, YAW)
T_W_slider = T_W_O @ T_O_slider
pusher_pos_3d_W = T_W_O @ pusher_pos_homog_O
pusher_velocity_3d_W = (T_W_O[:3, :3] @ pusher_velocity_3d_O)
# slider
slider_pose_dict = transform_utils.matrix_to_dict(T_W_slider)
q_slider = np.concatenate(
[slider_pose_dict['quaternion'], slider_pose_dict['position']])
q_pusher = pusher_pos_3d_W[:2]
action = pusher_velocity_3d_W[:2]
action_sequence = None
if False:
angle_1 = np.deg2rad(0)
action_1 = np.array([np.cos(angle_1), np.sin(angle_1)])
angle_2 = np.deg2rad(-20)
action_2 = np.array([np.cos(angle_2), np.sin(angle_2)])
# expand it to be an action sequence
# [N, 2]
action_sequence_1 = PUSHER_VELOCITY * torch.Tensor(action_1).unsqueeze(
0).expand([N // 2, -1])
action_sequence_2 = PUSHER_VELOCITY * torch.Tensor(action_2).unsqueeze(
0).expand([N // 2, -1])
action_sequence = torch.cat((action_sequence_1, action_sequence_2),
dim=0)
print("action_sequence.shape", action_sequence.shape)
else:
action_sequence = torch.Tensor(action).unsqueeze(0).expand([N, -1])
return {
'T_W_O': T_W_O,
'pusher_pos_3d': pusher_pos_3d_W,
'T_W_slider': T_W_slider,
'q_slider': q_slider,
'q_pusher': q_pusher,
'action': action,
'action_sequence': action_sequence,
}
def drake_position_to_pose(q):
assert len(q) == 7 # [quat, pos]
quat = q[0:4]
pos = q[4:]
return transform_utils.transform_from_pose(pos, quat)
def slider_pose_from_observation(
obs, # the raw observation
): # 4 x 4 homogeneous transform
pos = obs['slider']['position']['translation']
quat = obs['slider']['position']['quaternion']
T_W_obj = transform_utils.transform_from_pose(pos, quat)
return T_W_obj
def func(
obs_raw,
T=None, # 3 x 3 homogeneous transform to apply to data, used for data augmentation
data_augmentation=False,
**kwargs, # for compatibility
):
pos = obs_raw['slider']['position']['translation']
quat = obs_raw['slider']['position']['quaternion']
T_W_obj = transform_utils.transform_from_pose(pos, quat)
# transform points to world frame
pts_W = transform_utils.transform_points_3D(T_W_obj, pts_obj)
# apply additional transform it was passed in
# T is typically
if T is not None:
pts_W = transform_utils.transform_points_3D(T, pts_W)
# noise augmentation
if data_augmentation and (
config_data_aug
is not None) and config_data_aug["keypoints"]["augment"]:
# print("slider: ADDING NOISE")
noise = np.random.normal(
scale=config_data_aug["keypoints"]["std_dev"],
size=pts_W.shape)
pts_W += noise
return torch.Tensor(pts_W)