File size: 3,450 Bytes
9de9fbf |
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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import tensorflow as tf
import tensorflow_datasets as tfds
from data.utils import clean_task_instruction, quaternion_to_euler
def load_dataset():
builder = tfds.builder('robomimic_ph/square_ph_image')
builder.download_and_prepare()
ds = builder.as_dataset(split='train', shuffle_files=True)
return ds
def terminate_act_to_bool(terminate_act: tf.Tensor) -> tf.Tensor:
"""
Convert terminate action to a boolean, where True means terminate.
"""
return tf.where(tf.equal(terminate_act, tf.constant(0.0, dtype=tf.float32)),tf.constant(False),tf.constant(True))
def process_step(step: dict) -> dict:
"""
Unify the action format and clean the task instruction.
DO NOT use python list, use tf.TensorArray instead.
"""
# format refers to https://www.tensorflow.org/datasets/catalog/robomimic_mg
# Convert raw action to our action
eef = step['action']
step['action'] = {}
action = step['action']
action['terminate'] = step['is_terminal']
eef_delta_pos = eef[:3]
eef_ang = quaternion_to_euler(eef[3:])
# No base found
# Concatenate the action
arm_action = tf.concat([eef_delta_pos, eef_ang], axis=0)
action['arm_concat'] = arm_action
# Write the action format
action['format'] = tf.constant(
"eef_delta_pos_x,eef_delta_pos_y,eef_delta_pos_z,eef_delta_angle_roll,eef_delta_angle_pitch,eef_delta_angle_yaw")
# Convert raw state to our state
state = step['observation']
arm_joint_pos = state['robot0_joint_pos']
arm_joint_vel = state['robot0_joint_vel']
gripper_pos = state['robot0_gripper_qpos']
gripper_vel = state['robot0_gripper_qvel']
eef_pos = state['robot0_eef_pos']
eef_ang = quaternion_to_euler(state['robot0_eef_quat'])
state['arm_concat'] = tf.concat([arm_joint_pos, arm_joint_vel, gripper_pos,gripper_vel,eef_pos,eef_ang], axis=0)
# convert to tf32
state['arm_concat'] = tf.cast(state['arm_concat'], tf.float32)
# Write the state format
state['format'] = tf.constant(
"arm_joint_0_pos,arm_joint_1_pos,arm_joint_2_pos,arm_joint_3_pos,arm_joint_4_pos,arm_joint_5_pos,arm_joint_6_pos,arm_joint_0_vel,arm_joint_1_vel,arm_joint_2_vel,arm_joint_3_vel,arm_joint_4_vel,arm_joint_5_vel,arm_joint_6_vel,gripper_joint_0_pos,gripper_joint_1_pos,gripper_joint_0_vel,gripper_joint_1_vel,eef_pos_x,eef_pos_y,eef_pos_z,eef_angle_roll,eef_angle_pitch,eef_angle_yaw")
# Clean the task instruction
# Define the replacements (old, new) as a dictionary
replacements = {
'_': ' ',
'1f': ' ',
'4f': ' ',
'-': ' ',
'50': ' ',
'55': ' ',
'56': ' ',
}
# manual added by lbg
instr = "move the square across the cube"
instr = clean_task_instruction(instr, replacements)
step['observation']['natural_language_instruction'] = instr
return step
if __name__ == "__main__":
import tensorflow_datasets as tfds
from data.utils import dataset_to_path
DATASET_DIR = 'data/datasets/openx_embod'
DATASET_NAME = 'roboturk'
# Load the dataset
dataset = tfds.builder_from_directory(
builder_dir=dataset_to_path(
DATASET_NAME, DATASET_DIR))
dataset = dataset.as_dataset(split='all').take(1)
# Inspect the dataset
ze=tf.constant(0.0)
for episode in dataset:
for step in episode['steps']:
print(step)
break
|