def flatten_episode_agilex(episode: dict) -> tf.data.Dataset:
    """
    Flatten the episode to a list of steps.
    """
    episode_dict = episode['episode_dict']
    dataset_name = episode['dataset_name']
    
    json_content, states, masks, acts = generate_json_state(
        episode_dict, dataset_name
    )

    # Calculate the past_states for each step
    # Each step has a window of previous states with size ACTION_CHUNK_SIZE
    # Use the first state to pad the states
    # past_states will have shape (num_steps, ACTION_CHUNK_SIZE, state_dim)
    first_state = tf.expand_dims(states[0], axis=0)
    first_state = tf.repeat(first_state, ACTION_CHUNK_SIZE-1, axis=0)
    padded_states = tf.concat([first_state, states], axis=0)
    indices = tf.range(ACTION_CHUNK_SIZE, tf.shape(states)[0] + ACTION_CHUNK_SIZE)
    past_states = tf.map_fn(
        lambda i: padded_states[i - ACTION_CHUNK_SIZE:i],
        indices,
        dtype=tf.float32
    )
    states_time_mask = tf.ones([tf.shape(states)[0]], dtype=tf.bool)
    padded_states_time_mask = tf.pad(states_time_mask, [[ACTION_CHUNK_SIZE-1, 0]], "CONSTANT", constant_values=False)
    past_states_time_mask = tf.map_fn(
        lambda i: padded_states_time_mask[i - ACTION_CHUNK_SIZE:i],
        indices,
        dtype=tf.bool
    )

    # NOTE bg the future states shall be actions
    # Calculate the future_states for each step
    # Each step has a window of future states with size ACTION_CHUNK_SIZE
    # Use the last action to pad the states
    # future_states will have shape (num_steps, ACTION_CHUNK_SIZE, state_dim)
    last_act = tf.expand_dims(acts[-1], axis=0)
    last_act = tf.repeat(last_act, ACTION_CHUNK_SIZE, axis=0)
    padded_states = tf.concat([acts, last_act], axis=0)
    # indices = tf.range(1, tf.shape(states)[0] + 1)
    indices = tf.range(0, tf.shape(acts)[0]) # NOTE time 0 action = time 1 state
    future_states = tf.map_fn(
        lambda i: padded_states[i:i + ACTION_CHUNK_SIZE],
        indices,
        dtype=tf.float32
    )
    states_time_mask = tf.ones([tf.shape(acts)[0]], dtype=tf.bool)
    padded_states_time_mask = tf.pad(states_time_mask, [[0, ACTION_CHUNK_SIZE]], "CONSTANT", constant_values=False)
    future_states_time_mask = tf.map_fn(
        lambda i: padded_states_time_mask[i:i + ACTION_CHUNK_SIZE],
        indices,
        dtype=tf.bool
    )

    # Calculate the std and mean for state
    state_std = tf.math.reduce_std(states, axis=0, keepdims=True)
    state_std = tf.repeat(state_std, tf.shape(states)[0], axis=0)
    state_mean = tf.math.reduce_mean(states, axis=0, keepdims=True)
    state_mean = tf.repeat(state_mean, tf.shape(states)[0], axis=0)

    state_norm = tf.math.reduce_mean(
        tf.math.square(acts), axis=0, keepdims=True)
    state_norm = tf.math.sqrt(state_norm)
    state_norm = tf.repeat(state_norm, tf.shape(states)[0], axis=0)

    # Create a list of steps
    step_data = []
    for i in range(tf.shape(states)[0]):
        step_data.append({
            'step_id': episode['step_id'][i],
            'json_content': json_content,
            'state_chunk': past_states[i],
            'state_chunk_time_mask': past_states_time_mask[i],
            'action_chunk': future_states[i],
            'action_chunk_time_mask': future_states_time_mask[i],
            'state_vec_mask': masks[i],
            'past_frames_0': episode['past_frames_0'][i],
            'past_frames_0_time_mask': episode['past_frames_0_time_mask'][i],
            'past_frames_1': episode['past_frames_1'][i],
            'past_frames_1_time_mask': episode['past_frames_1_time_mask'][i],
            'past_frames_2': episode['past_frames_2'][i],
            'past_frames_2_time_mask': episode['past_frames_2_time_mask'][i],
            'past_frames_3': episode['past_frames_3'][i],
            'past_frames_3_time_mask': episode['past_frames_3_time_mask'][i],
            'state_std': state_std[i],
            'state_mean': state_mean[i],
            'state_norm': state_norm[i],
        })
    
    return step_data

    Flatten the episode to a list of steps.
    
policy/RDT/data/episode_transform.py