File size: 2,603 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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import tensorflow as tf
from data.utils import clean_task_instruction, euler_to_quaternion
def terminate_act_to_bool(terminate_act: tf.Tensor) -> tf.Tensor:
"""
Convert terminate action to a boolean, where True means terminate.
"""
return tf.equal(terminate_act, tf.constant(1.0, dtype=tf.float32))
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 = step['action']
action['terminate'] = terminate_act_to_bool(action['terminate_episode'])
# Multiplied by 3 Hz to get units m/s and rad/s
eef_delta_pos = action['world_vector'] * 3
eef_ang = action['rotation_delta'] * 3
# Origin: [-0.28, 0.96]: open, close
# 1-Origin: [0.04, 1.28]: close, open
grip_open = 1 - action['gripper_closedness_action']
# base_delta_pos = action['base_displacement_vector']
# base_delta_ang = action['base_displacement_vertical_rotation']
# Concatenate the action
arm_action = tf.concat([eef_delta_pos, eef_ang, grip_open], axis=0)
action['arm_concat'] = arm_action
# base_action = tf.concat([base_delta_pos, base_delta_ang], axis=0)
# action['base_concat'] = base_action
# Write the action format
action['format'] = tf.constant(
"eef_vel_x,eef_vel_y,eef_vel_z,eef_angular_vel_roll,eef_angular_vel_pitch,eef_angular_vel_yaw,gripper_open")
# Convert raw state to our state
# state = step['observation']
# eef_pos = state['base_pose_tool_reached'][:3]
# eef_ang = quaternion_to_euler(state['base_pose_tool_reached'][3:])
# grip_open = 1 - state['gripper_closed']
# create empty tensor
# state['arm_concat'] = tf.constant([0, 0, 0, 0, 0, 0], dtype=tf.float32)
# Write the state format
# state['format'] = tf.constant(
# "")
# Define the task instruction
step['observation']['natural_language_instruction'] = tf.constant(
"Open the cabinet door.")
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 = 'nyu_door_opening_surprising_effectiveness'
# Load the dataset
dataset = tfds.builder_from_directory(
builder_dir=dataset_to_path(
DATASET_NAME, DATASET_DIR))
dataset = dataset.as_dataset(split='all')
# Inspect the dataset
for episode in dataset.take(1):
for step in episode['steps']:
print(step)
|