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)