def process_episode(epsd: dict, dataset_name: str, 
                    image_keys: list, image_mask: list) -> dict:
    """
    Process an episode to extract the frames and the json content.
    """
    # Frames of each camera
    # Ugly code due to tf's poor compatibility
    frames_0 = tf.TensorArray(dtype=tf.uint8, size=0, dynamic_size=True)
    frames_1 = tf.TensorArray(dtype=tf.uint8, size=0, dynamic_size=True)
    frames_2 = tf.TensorArray(dtype=tf.uint8, size=0, dynamic_size=True)
    frames_3 = tf.TensorArray(dtype=tf.uint8, size=0, dynamic_size=True)
    # Traverse the episode to collect...
    for step in iter(epsd['steps']):
        # Parse the image
        frames_0 = frames_0.write(frames_0.size(), 
            tf.cond(
                tf.equal(image_mask[0], 1),
                lambda: step['observation'][image_keys[0]],
                lambda: tf.zeros([0, 0, 0], dtype=tf.uint8)
            ))
        # Very ugly code due to tf's poor compatibility
        frames_1 = frames_1.write(frames_1.size(), 
            tf.cond(
                tf.equal(image_mask[1], 1),
                lambda: step['observation'][image_keys[1]],
                lambda: tf.zeros([0, 0, 0], dtype=tf.uint8)
            ))
        frames_2 = frames_2.write(frames_2.size(), 
            tf.cond(
                tf.equal(image_mask[2], 1),
                lambda: step['observation'][image_keys[2]],
                lambda: tf.zeros([0, 0, 0], dtype=tf.uint8)
            ))
        frames_3 = frames_3.write(frames_3.size(),
            tf.cond(
                tf.equal(image_mask[3], 1),
                lambda: step['observation'][image_keys[3]],
                lambda: tf.zeros([0, 0, 0], dtype=tf.uint8)
            ))
    

    # Calculate the past_frames_0 for each step
    # Each step has a window of previous frames with size IMG_HISTORY_SIZE
    # Use the first state to pad the frames
    # past_frames_0 will have shape (num_steps, IMG_HISTORY_SIZE, height, width, channels)
    frames_0 = frames_0.stack()
    first_frame = tf.expand_dims(frames_0[0], axis=0)
    first_frame = tf.repeat(first_frame, IMG_HISTORY_SIZE-1, axis=0)
    padded_frames_0 = tf.concat([first_frame, frames_0], axis=0)
    indices = tf.range(IMG_HISTORY_SIZE, tf.shape(frames_0)[0] + IMG_HISTORY_SIZE)
    past_frames_0 = tf.map_fn(
        lambda i: padded_frames_0[i - IMG_HISTORY_SIZE:i],
        indices,
        dtype=tf.uint8
    )
    frames_0_time_mask = tf.ones([tf.shape(frames_0)[0]], dtype=tf.bool)
    padded_frames_0_time_mask = tf.pad(frames_0_time_mask, [[IMG_HISTORY_SIZE-1, 0]], "CONSTANT", constant_values=False)
    past_frames_0_time_mask = tf.map_fn(
        lambda i: padded_frames_0_time_mask[i - IMG_HISTORY_SIZE:i],
        indices,
        dtype=tf.bool
    )

    # For past_frames_1
    frames_1 = frames_1.stack()
    first_frame = tf.expand_dims(frames_1[0], axis=0)
    first_frame = tf.repeat(first_frame, IMG_HISTORY_SIZE-1, axis=0)
    padded_frames_1 = tf.concat([first_frame, frames_1], axis=0)
    indices = tf.range(IMG_HISTORY_SIZE, tf.shape(frames_1)[0] + IMG_HISTORY_SIZE)
    past_frames_1 = tf.map_fn(
        lambda i: padded_frames_1[i - IMG_HISTORY_SIZE:i],
        indices,
        dtype=tf.uint8
    )
    frames_1_time_mask = tf.ones([tf.shape(frames_1)[0]], dtype=tf.bool)
    padded_frames_1_time_mask = tf.pad(frames_1_time_mask, [[IMG_HISTORY_SIZE-1, 0]], "CONSTANT", constant_values=False)
    past_frames_1_time_mask = tf.map_fn(
        lambda i: padded_frames_1_time_mask[i - IMG_HISTORY_SIZE:i],
        indices,
        dtype=tf.bool
    )

    # For past_frames_2
    frames_2 = frames_2.stack()
    first_frame = tf.expand_dims(frames_2[0], axis=0)
    first_frame = tf.repeat(first_frame, IMG_HISTORY_SIZE-1, axis=0)
    padded_frames_2 = tf.concat([first_frame, frames_2], axis=0)
    indices = tf.range(IMG_HISTORY_SIZE, tf.shape(frames_2)[0] + IMG_HISTORY_SIZE)
    past_frames_2 = tf.map_fn(
        lambda i: padded_frames_2[i - IMG_HISTORY_SIZE:i],
        indices,
        dtype=tf.uint8
    )
    frames_2_time_mask = tf.ones([tf.shape(frames_2)[0]], dtype=tf.bool)
    padded_frames_2_time_mask = tf.pad(frames_2_time_mask, [[IMG_HISTORY_SIZE-1, 0]], "CONSTANT", constant_values=False)
    past_frames_2_time_mask = tf.map_fn(
        lambda i: padded_frames_2_time_mask[i - IMG_HISTORY_SIZE:i],
        indices,
        dtype=tf.bool
    )

    # For past_frames_3
    frames_3 = frames_3.stack()
    first_frame = tf.expand_dims(frames_3[0], axis=0)
    first_frame = tf.repeat(first_frame, IMG_HISTORY_SIZE-1, axis=0)
    padded_frames_3 = tf.concat([first_frame, frames_3], axis=0)
    indices = tf.range(IMG_HISTORY_SIZE, tf.shape(frames_3)[0] + IMG_HISTORY_SIZE)
    past_frames_3 = tf.map_fn(
        lambda i: padded_frames_3[i - IMG_HISTORY_SIZE:i],
        indices,
        dtype=tf.uint8
    )
    frames_3_time_mask = tf.ones([tf.shape(frames_3)[0]], dtype=tf.bool)
    padded_frames_3_time_mask = tf.pad(frames_3_time_mask, [[IMG_HISTORY_SIZE-1, 0]], "CONSTANT", constant_values=False)
    past_frames_3_time_mask = tf.map_fn(
        lambda i: padded_frames_3_time_mask[i - IMG_HISTORY_SIZE:i],
        indices,
        dtype=tf.bool
    )

    # Creat the ids for each step
    step_id = tf.range(0, tf.shape(frames_0)[0])

    return {
        'dataset_name': dataset_name,
        'episode_dict': epsd,
        'step_id': step_id,
        'past_frames_0': past_frames_0,
        'past_frames_0_time_mask': past_frames_0_time_mask,
        'past_frames_1': past_frames_1,
        'past_frames_1_time_mask': past_frames_1_time_mask,
        'past_frames_2': past_frames_2,
        'past_frames_2_time_mask': past_frames_2_time_mask,
        'past_frames_3': past_frames_3,
        'past_frames_3_time_mask': past_frames_3_time_mask,
    }

    Process an episode to extract the frames and the json content.
    
policy/RDT/data/episode_transform.py