File size: 2,008 Bytes
16628c8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import tensorflow as tf
from data.utils import clean_task_instruction, euler_to_rotation_matrix, rotation_matrix_to_ortho6d
def process_step(step: dict) -> dict:
"""
Unify the action format and clean the task instruction.
DO NOT use python list, use tf.TensorArray instead.
"""
# Convert raw action to our action
action_dict = step['action']
# Concatenate the action
step['action'] = {}
action = step['action']
action['arm_concat'] = action_dict['ee_6d_pos']
# Write the action format
action['format'] = tf.constant(
"left_eef_pos_x,left_eef_pos_y,left_eef_pos_z,left_eef_angle_0,left_eef_angle_1,left_eef_angle_2,left_eef_angle_3,left_eef_angle_4,left_eef_angle_5,left_gripper_open,right_eef_pos_x,right_eef_pos_y,right_eef_pos_z,right_eef_angle_0,right_eef_angle_1,right_eef_angle_2,right_eef_angle_3,right_eef_angle_4,right_eef_angle_5,right_gripper_open"
)
# Convert raw state to our state
# Robot state
state_dict = step['observation']['state']
state = {}
state['arm_concat'] = state_dict
# Write the state format
state['format'] = tf.constant(
"left_eef_pos_x,left_eef_pos_y,left_eef_pos_z,left_eef_angle_0,left_eef_angle_1,left_eef_angle_2,left_eef_angle_3,left_eef_angle_4,left_eef_angle_5,left_gripper_open,right_eef_pos_x,right_eef_pos_y,right_eef_pos_z,right_eef_angle_0,right_eef_angle_1,right_eef_angle_2,right_eef_angle_3,right_eef_angle_4,right_eef_angle_5,right_gripper_open"
)
# Clean the task instruction
# Define the replacements (old, new) as a dictionary
replacements = {
'_': ' ',
'1f': ' ',
'4f': ' ',
'-': ' ',
'50': ' ',
'55': ' ',
'56': ' ',
}
instr = step['language_instruction']
# instr = clean_task_instruction(instr, replacements)
step['observation'] = state
step['observation']['natural_language_instruction'] = instr
return step
if __name__ == "__main__":
pass
|