156 lines
6.5 KiB
Python
156 lines
6.5 KiB
Python
import pathlib
|
|
|
|
### Task parameters
|
|
DATA_DIR = str(pathlib.Path(__file__).parent.resolve() / 'data')
|
|
SIM_TASK_CONFIGS = {
|
|
'sim_transfer_cube_scripted':{
|
|
'dataset_dir': DATA_DIR + '/sim_transfer_cube_scripted',
|
|
'num_episodes': 50,
|
|
'episode_len': 400,
|
|
'camera_names': ['top']
|
|
},
|
|
|
|
'sim_transfer_cube_human':{
|
|
'dataset_dir': DATA_DIR + '/sim_transfer_cube_human',
|
|
'num_episodes': 50,
|
|
'episode_len': 400,
|
|
'camera_names': ['top']
|
|
},
|
|
|
|
'sim_insertion_scripted': {
|
|
'dataset_dir': DATA_DIR + '/sim_insertion_scripted',
|
|
'num_episodes': 50,
|
|
'episode_len': 400,
|
|
'camera_names': ['top']
|
|
},
|
|
|
|
'sim_insertion_human': {
|
|
'dataset_dir': DATA_DIR + '/sim_insertion_human',
|
|
'num_episodes': 50,
|
|
'episode_len': 500,
|
|
'camera_names': ['top']
|
|
},
|
|
}
|
|
|
|
ENDOSCOPE_TASK_CONFIGS = {
|
|
'endoscope_default': {
|
|
'dataset_dir': DATA_DIR + '/endoscope_default',
|
|
'num_episodes': 50,
|
|
'episode_len': 400,
|
|
'camera_names': ['top'],
|
|
'state_dim': 2,
|
|
'action_dim': 2,
|
|
'real_action_t_minus_1': False,
|
|
'use_text_instruction': True,
|
|
'instruction_mode': 'timestep-level',
|
|
'use_cached_text_features': True,
|
|
'text_encoder_type': 'distilbert',
|
|
'text_feature_dim': 768,
|
|
'text_fusion_type': 'concat_transformer_input',
|
|
'freeze_text_encoder': True,
|
|
'text_max_length': 32,
|
|
'text_tokenizer_name': 'distilbert-base-uncased',
|
|
},
|
|
'endoscope_follow': {
|
|
'dataset_dir': DATA_DIR + '/follow',
|
|
'num_episodes': 3,
|
|
'episode_len': 400,
|
|
'camera_names': ['top'],
|
|
'state_dim': 2,
|
|
'action_dim': 2,
|
|
'real_action_t_minus_1': False,
|
|
'use_text_instruction': True,
|
|
'instruction_mode': 'timestep-level',
|
|
'use_cached_text_features': True,
|
|
'text_encoder_type': 'distilbert',
|
|
'text_feature_dim': 768,
|
|
'text_fusion_type': 'concat_transformer_input',
|
|
'freeze_text_encoder': True,
|
|
'text_max_length': 32,
|
|
'text_tokenizer_name': 'distilbert-base-uncased',
|
|
},
|
|
'endoscope_both_no_text': {
|
|
'dataset_dir': DATA_DIR + '/both-no-text',
|
|
'num_episodes': 3,
|
|
'episode_len': 400,
|
|
'camera_names': ['top'],
|
|
'state_dim': 2,
|
|
'action_dim': 2,
|
|
'real_action_t_minus_1': False,
|
|
'use_text_instruction': False,
|
|
},
|
|
'endoscope_sanity_check': {
|
|
'dataset_dir': DATA_DIR + '/sanity-check',
|
|
'num_episodes': 3,
|
|
'episode_len': 400,
|
|
'camera_names': ['top'],
|
|
'state_dim': 2,
|
|
'action_dim': 2,
|
|
'real_action_t_minus_1': False,
|
|
'use_text_instruction': False,
|
|
},
|
|
'endoscope_cannulation_no_text': {
|
|
'dataset_dir': DATA_DIR + '/cannulation-no-text',
|
|
'num_episodes': 3,
|
|
'episode_len': 400,
|
|
'camera_names': ['top'],
|
|
'state_dim': 2,
|
|
'action_dim': 2,
|
|
'real_action_t_minus_1': False,
|
|
'use_text_instruction': False,
|
|
},
|
|
'endoscope_follow_no_text': {
|
|
'dataset_dir': DATA_DIR + '/follow-no-text',
|
|
'num_episodes': 3,
|
|
'episode_len': 400,
|
|
'camera_names': ['top'],
|
|
'state_dim': 2,
|
|
'action_dim': 2,
|
|
'real_action_t_minus_1': False,
|
|
'use_text_instruction': False,
|
|
},
|
|
}
|
|
|
|
### Simulation envs fixed constants
|
|
DT = 0.02
|
|
JOINT_NAMES = ["waist", "shoulder", "elbow", "forearm_roll", "wrist_angle", "wrist_rotate"]
|
|
START_ARM_POSE = [0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239, 0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239]
|
|
|
|
XML_DIR = str(pathlib.Path(__file__).parent.resolve()) + '/assets/' # note: absolute path
|
|
|
|
# Left finger position limits (qpos[7]), right_finger = -1 * left_finger
|
|
MASTER_GRIPPER_POSITION_OPEN = 0.02417
|
|
MASTER_GRIPPER_POSITION_CLOSE = 0.01244
|
|
PUPPET_GRIPPER_POSITION_OPEN = 0.05800
|
|
PUPPET_GRIPPER_POSITION_CLOSE = 0.01844
|
|
|
|
# Gripper joint limits (qpos[6])
|
|
MASTER_GRIPPER_JOINT_OPEN = 0.3083
|
|
MASTER_GRIPPER_JOINT_CLOSE = -0.6842
|
|
PUPPET_GRIPPER_JOINT_OPEN = 1.4910
|
|
PUPPET_GRIPPER_JOINT_CLOSE = -0.6213
|
|
|
|
############################ Helper functions ############################
|
|
|
|
MASTER_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_POSITION_CLOSE) / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
|
|
PUPPET_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_POSITION_CLOSE) / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)
|
|
MASTER_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) + MASTER_GRIPPER_POSITION_CLOSE
|
|
PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + PUPPET_GRIPPER_POSITION_CLOSE
|
|
MASTER2PUPPET_POSITION_FN = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(MASTER_GRIPPER_POSITION_NORMALIZE_FN(x))
|
|
|
|
MASTER_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE)
|
|
PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE)
|
|
MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
|
|
PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
|
|
MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x))
|
|
|
|
MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
|
|
PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)
|
|
|
|
MASTER_POS2JOINT = lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
|
|
MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN((x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE))
|
|
PUPPET_POS2JOINT = lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
|
|
PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN((x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE))
|
|
|
|
MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE)/2
|