|
|
from equi_diffpo.model.common.normalizer import SingleFieldLinearNormalizer |
|
|
from equi_diffpo.common.pytorch_util import dict_apply, dict_apply_reduce, dict_apply_split |
|
|
import numpy as np |
|
|
|
|
|
|
|
|
def get_range_normalizer_from_stat(stat, output_max=1, output_min=-1, range_eps=1e-7): |
|
|
|
|
|
input_max = stat['max'] |
|
|
input_min = stat['min'] |
|
|
input_range = input_max - input_min |
|
|
ignore_dim = input_range < range_eps |
|
|
input_range[ignore_dim] = output_max - output_min |
|
|
scale = (output_max - output_min) / input_range |
|
|
offset = output_min - scale * input_min |
|
|
offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] |
|
|
|
|
|
return SingleFieldLinearNormalizer.create_manual( |
|
|
scale=scale, |
|
|
offset=offset, |
|
|
input_stats_dict=stat |
|
|
) |
|
|
|
|
|
def get_range_symmetric_normalizer_from_stat(stat, output_max=1, output_min=-1, range_eps=1e-7): |
|
|
|
|
|
input_max = stat['max'] |
|
|
input_min = stat['min'] |
|
|
abs_max = np.max([np.abs(stat['max'][:2]), np.abs(stat['min'][:2])]) |
|
|
input_max[:2] = abs_max |
|
|
input_min[:2] = -abs_max |
|
|
input_range = input_max - input_min |
|
|
ignore_dim = input_range < range_eps |
|
|
input_range[ignore_dim] = output_max - output_min |
|
|
scale = (output_max - output_min) / input_range |
|
|
offset = output_min - scale * input_min |
|
|
offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] |
|
|
|
|
|
return SingleFieldLinearNormalizer.create_manual( |
|
|
scale=scale, |
|
|
offset=offset, |
|
|
input_stats_dict=stat |
|
|
) |
|
|
|
|
|
def get_voxel_identity_normalizer(): |
|
|
scale = np.array([1], dtype=np.float32) |
|
|
offset = np.array([0], dtype=np.float32) |
|
|
stat = { |
|
|
'min': np.array([0], dtype=np.float32), |
|
|
'max': np.array([1], dtype=np.float32), |
|
|
'mean': np.array([0.5], dtype=np.float32), |
|
|
'std': np.array([np.sqrt(1/12)], dtype=np.float32) |
|
|
} |
|
|
return SingleFieldLinearNormalizer.create_manual( |
|
|
scale=scale, |
|
|
offset=offset, |
|
|
input_stats_dict=stat |
|
|
) |
|
|
|
|
|
def get_image_range_normalizer(): |
|
|
scale = np.array([2], dtype=np.float32) |
|
|
offset = np.array([-1], dtype=np.float32) |
|
|
stat = { |
|
|
'min': np.array([0], dtype=np.float32), |
|
|
'max': np.array([1], dtype=np.float32), |
|
|
'mean': np.array([0.5], dtype=np.float32), |
|
|
'std': np.array([np.sqrt(1/12)], dtype=np.float32) |
|
|
} |
|
|
return SingleFieldLinearNormalizer.create_manual( |
|
|
scale=scale, |
|
|
offset=offset, |
|
|
input_stats_dict=stat |
|
|
) |
|
|
|
|
|
def get_identity_normalizer_from_stat(stat): |
|
|
scale = np.ones_like(stat['min']) |
|
|
offset = np.zeros_like(stat['min']) |
|
|
return SingleFieldLinearNormalizer.create_manual( |
|
|
scale=scale, |
|
|
offset=offset, |
|
|
input_stats_dict=stat |
|
|
) |
|
|
|
|
|
def robomimic_abs_action_normalizer_from_stat(stat, rotation_transformer): |
|
|
result = dict_apply_split( |
|
|
stat, lambda x: { |
|
|
'pos': x[...,:3], |
|
|
'rot': x[...,3:6], |
|
|
'gripper': x[...,6:] |
|
|
}) |
|
|
|
|
|
def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7): |
|
|
|
|
|
input_max = stat['max'] |
|
|
input_min = stat['min'] |
|
|
input_range = input_max - input_min |
|
|
ignore_dim = input_range < range_eps |
|
|
input_range[ignore_dim] = output_max - output_min |
|
|
scale = (output_max - output_min) / input_range |
|
|
offset = output_min - scale * input_min |
|
|
offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] |
|
|
|
|
|
return {'scale': scale, 'offset': offset}, stat |
|
|
|
|
|
def get_rot_param_info(stat): |
|
|
example = rotation_transformer.forward(stat['mean']) |
|
|
scale = np.ones_like(example) |
|
|
offset = np.zeros_like(example) |
|
|
info = { |
|
|
'max': np.ones_like(example), |
|
|
'min': np.full_like(example, -1), |
|
|
'mean': np.zeros_like(example), |
|
|
'std': np.ones_like(example) |
|
|
} |
|
|
return {'scale': scale, 'offset': offset}, info |
|
|
|
|
|
def get_gripper_param_info(stat): |
|
|
example = stat['max'] |
|
|
scale = np.ones_like(example) |
|
|
offset = np.zeros_like(example) |
|
|
info = { |
|
|
'max': np.ones_like(example), |
|
|
'min': np.full_like(example, -1), |
|
|
'mean': np.zeros_like(example), |
|
|
'std': np.ones_like(example) |
|
|
} |
|
|
return {'scale': scale, 'offset': offset}, info |
|
|
|
|
|
pos_param, pos_info = get_pos_param_info(result['pos']) |
|
|
rot_param, rot_info = get_rot_param_info(result['rot']) |
|
|
gripper_param, gripper_info = get_gripper_param_info(result['gripper']) |
|
|
|
|
|
param = dict_apply_reduce( |
|
|
[pos_param, rot_param, gripper_param], |
|
|
lambda x: np.concatenate(x,axis=-1)) |
|
|
info = dict_apply_reduce( |
|
|
[pos_info, rot_info, gripper_info], |
|
|
lambda x: np.concatenate(x,axis=-1)) |
|
|
|
|
|
return SingleFieldLinearNormalizer.create_manual( |
|
|
scale=param['scale'], |
|
|
offset=param['offset'], |
|
|
input_stats_dict=info |
|
|
) |
|
|
|
|
|
|
|
|
def robomimic_abs_action_only_normalizer_from_stat(stat): |
|
|
result = dict_apply_split( |
|
|
stat, lambda x: { |
|
|
'pos': x[...,:3], |
|
|
'other': x[...,3:] |
|
|
}) |
|
|
|
|
|
def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7): |
|
|
|
|
|
input_max = stat['max'] |
|
|
input_min = stat['min'] |
|
|
input_range = input_max - input_min |
|
|
ignore_dim = input_range < range_eps |
|
|
input_range[ignore_dim] = output_max - output_min |
|
|
scale = (output_max - output_min) / input_range |
|
|
offset = output_min - scale * input_min |
|
|
offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] |
|
|
|
|
|
return {'scale': scale, 'offset': offset}, stat |
|
|
|
|
|
|
|
|
def get_other_param_info(stat): |
|
|
example = stat['max'] |
|
|
scale = np.ones_like(example) |
|
|
offset = np.zeros_like(example) |
|
|
info = { |
|
|
'max': np.ones_like(example), |
|
|
'min': np.full_like(example, -1), |
|
|
'mean': np.zeros_like(example), |
|
|
'std': np.ones_like(example) |
|
|
} |
|
|
return {'scale': scale, 'offset': offset}, info |
|
|
|
|
|
pos_param, pos_info = get_pos_param_info(result['pos']) |
|
|
other_param, other_info = get_other_param_info(result['other']) |
|
|
|
|
|
param = dict_apply_reduce( |
|
|
[pos_param, other_param], |
|
|
lambda x: np.concatenate(x,axis=-1)) |
|
|
info = dict_apply_reduce( |
|
|
[pos_info, other_info], |
|
|
lambda x: np.concatenate(x,axis=-1)) |
|
|
|
|
|
return SingleFieldLinearNormalizer.create_manual( |
|
|
scale=param['scale'], |
|
|
offset=param['offset'], |
|
|
input_stats_dict=info |
|
|
) |
|
|
|
|
|
|
|
|
def robomimic_abs_action_only_symmetric_normalizer_from_stat(stat): |
|
|
result = dict_apply_split( |
|
|
stat, lambda x: { |
|
|
'pos': x[...,:3], |
|
|
'other': x[...,3:] |
|
|
}) |
|
|
|
|
|
def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7): |
|
|
|
|
|
input_max = stat['max'] |
|
|
input_min = stat['min'] |
|
|
abs_max = np.max([np.abs(stat['max'][:2]), np.abs(stat['min'][:2])]) |
|
|
input_max[:2] = abs_max |
|
|
input_min[:2] = -abs_max |
|
|
input_range = input_max - input_min |
|
|
ignore_dim = input_range < range_eps |
|
|
input_range[ignore_dim] = output_max - output_min |
|
|
scale = (output_max - output_min) / input_range |
|
|
offset = output_min - scale * input_min |
|
|
offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] |
|
|
|
|
|
return {'scale': scale, 'offset': offset}, stat |
|
|
|
|
|
|
|
|
def get_other_param_info(stat): |
|
|
example = stat['max'] |
|
|
scale = np.ones_like(example) |
|
|
offset = np.zeros_like(example) |
|
|
info = { |
|
|
'max': np.ones_like(example), |
|
|
'min': np.full_like(example, -1), |
|
|
'mean': np.zeros_like(example), |
|
|
'std': np.ones_like(example) |
|
|
} |
|
|
return {'scale': scale, 'offset': offset}, info |
|
|
|
|
|
pos_param, pos_info = get_pos_param_info(result['pos']) |
|
|
other_param, other_info = get_other_param_info(result['other']) |
|
|
|
|
|
param = dict_apply_reduce( |
|
|
[pos_param, other_param], |
|
|
lambda x: np.concatenate(x,axis=-1)) |
|
|
info = dict_apply_reduce( |
|
|
[pos_info, other_info], |
|
|
lambda x: np.concatenate(x,axis=-1)) |
|
|
|
|
|
return SingleFieldLinearNormalizer.create_manual( |
|
|
scale=param['scale'], |
|
|
offset=param['offset'], |
|
|
input_stats_dict=info |
|
|
) |
|
|
|
|
|
|
|
|
def robomimic_abs_action_only_dual_arm_normalizer_from_stat(stat): |
|
|
Da = stat['max'].shape[-1] |
|
|
Dah = Da // 2 |
|
|
result = dict_apply_split( |
|
|
stat, lambda x: { |
|
|
'pos0': x[...,:3], |
|
|
'other0': x[...,3:Dah], |
|
|
'pos1': x[...,Dah:Dah+3], |
|
|
'other1': x[...,Dah+3:] |
|
|
}) |
|
|
|
|
|
def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7): |
|
|
|
|
|
input_max = stat['max'] |
|
|
input_min = stat['min'] |
|
|
input_range = input_max - input_min |
|
|
ignore_dim = input_range < range_eps |
|
|
input_range[ignore_dim] = output_max - output_min |
|
|
scale = (output_max - output_min) / input_range |
|
|
offset = output_min - scale * input_min |
|
|
offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] |
|
|
|
|
|
return {'scale': scale, 'offset': offset}, stat |
|
|
|
|
|
|
|
|
def get_other_param_info(stat): |
|
|
example = stat['max'] |
|
|
scale = np.ones_like(example) |
|
|
offset = np.zeros_like(example) |
|
|
info = { |
|
|
'max': np.ones_like(example), |
|
|
'min': np.full_like(example, -1), |
|
|
'mean': np.zeros_like(example), |
|
|
'std': np.ones_like(example) |
|
|
} |
|
|
return {'scale': scale, 'offset': offset}, info |
|
|
|
|
|
pos0_param, pos0_info = get_pos_param_info(result['pos0']) |
|
|
pos1_param, pos1_info = get_pos_param_info(result['pos1']) |
|
|
other0_param, other0_info = get_other_param_info(result['other0']) |
|
|
other1_param, other1_info = get_other_param_info(result['other1']) |
|
|
|
|
|
param = dict_apply_reduce( |
|
|
[pos0_param, other0_param, pos1_param, other1_param], |
|
|
lambda x: np.concatenate(x,axis=-1)) |
|
|
info = dict_apply_reduce( |
|
|
[pos0_info, other0_info, pos1_info, other1_info], |
|
|
lambda x: np.concatenate(x,axis=-1)) |
|
|
|
|
|
return SingleFieldLinearNormalizer.create_manual( |
|
|
scale=param['scale'], |
|
|
offset=param['offset'], |
|
|
input_stats_dict=info |
|
|
) |
|
|
|
|
|
|
|
|
def array_to_stats(arr: np.ndarray): |
|
|
stat = { |
|
|
'min': np.min(arr, axis=0), |
|
|
'max': np.max(arr, axis=0), |
|
|
'mean': np.mean(arr, axis=0), |
|
|
'std': np.std(arr, axis=0) |
|
|
} |
|
|
return stat |
|
|
|
