1t5-v1_1-xxl：
2bos://cce-ai-datasets/cce-ai-datasets.bj.bcebos.com/robotics-diffusion-transformer/rdt-1b/google/t5-v1_1-xxl
3
4siglip：
5bos://cce-ai-datasets/cce-ai-datasets.bj.bcebos.com/robotics-diffusion-transformer/rdt-1b/google/siglip-so400m-patch14-384

1bos://cce-ai-datasets/cce-ai-datasets.bj.bcebos.com/robotics-diffusion-transformer/rdt-1b/data/austin_sailor_dataset_converted_externally_to_rlds

1import tensorflow as tf
2
3from data.utils import clean_task_instruction, euler_to_quaternion, rotation_matrix_to_ortho6d
4
5def process_step(step: dict) -> dict:
6    """
7    Unify the action format and clean the task instruction.
8
9    DO NOT use python list, use tf.TensorArray instead.
10    """
11    # Convert raw action to our action
12
13    origin_action = step['action']
14    step['action']={}
15    action=step['action']
16    action['terminate'] = step['is_terminal']
17
18    eef_delta_pos = origin_action[:3]
19    eef_ang=origin_action[3:6]
20    eef_ang = euler_to_quaternion(eef_ang)
21    # gripper_open: -1-open, 1-closed
22    grip_open=tf.where(tf.equal(origin_action[6:],tf.constant(-1.0)),tf.constant(1.0),tf.constant(0.0))
23
24    # No base found
25
26    # Concatenate the action
27    action['arm_concat'] = tf.concat([eef_delta_pos,eef_ang,grip_open],axis=0)
28
29    # Write the action format
30    action['format'] = tf.constant(
31        "eef_delta_pos_x,eef_delta_pos_y,eef_delta_pos_z,eef_delta_angle_x,eef_delta_angle_y,eef_delta_angle_z,eef_delta_angle_w,gripper_open")
32
33    # Convert raw state to our state
34    state = step['observation']
35    # Concatenate the state
36    eef_mat = tf.transpose(tf.reshape(state['state'][8:], (4, 4)))
37    eef_pos = eef_mat[:3, 3]
38    rotaion_matrix = eef_mat[:3, :3]
39    eef_ang = rotation_matrix_to_ortho6d(rotaion_matrix)
40    joint_pos = state['state'][:7]
41    grip_open = state['state'][7:8] * 12.5 # rescale to [0, 1]
42    state['arm_concat'] = tf.concat([joint_pos,grip_open,eef_pos,eef_ang],axis=0)
43
44    # Write the state format
45    state['format'] = tf.constant(
46        "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,gripper_joint_0_pos,eef_pos_x,eef_pos_y,eef_pos_z,eef_angle_0,eef_angle_1,eef_angle_2,eef_angle_3,eef_angle_4,eef_angle_5")
47
48    # Clean the task instruction
49    # Define the replacements (old, new) as a dictionary
50    replacements = {
51        '_': ' ',
52        '1f': ' ',
53        '4f': ' ',
54        '-': ' ',
55        '50': ' ',
56        '55': ' ',
57        '56': ' ',
58        
59    }
60    instr = step['language_instruction']
61    instr = clean_task_instruction(instr, replacements)
62    step['observation']['natural_language_instruction'] = instr
63
64    return step
65
66
67if __name__ == "__main__":
68    import tensorflow_datasets as tfds
69    from data.utils import dataset_to_path
70
71    DATASET_DIR = 'data/datasets/openx_embod'
72    DATASET_NAME = 'rdt_data'
73    # Load the dataset
74    dataset = tfds.builder_from_directory(
75        builder_dir=dataset_to_path(
76            DATASET_NAME, DATASET_DIR))
77    dataset = dataset.as_dataset(split='all')
78
79    # Inspect the dataset
80    for episode in dataset:
81        for step in episode['steps']:
82            print(step)

1bos://cce-ai-datasets/cce-ai-datasets.bj.bcebos.com/robotics-diffusion-transformer/rdt-1b/data/rdt_ft_data

1bos://cce-ai-datasets/cce-ai-datasets.bj.bcebos.com/robotics-diffusion-transformer/rdt-1b/ckpt/rdt-1b