Data loader normalization fn issue

[fredshentu]

https://github.com/google-research/language-table/blob/main/language_table/train/input_pipeline_rlds.py#L256 See code here I don't think the tf.nest.flatten will preserve the keyvalue order and in fact, if you run the call this dataset creation method input_pipeline_rlds.create_datasets the dataloader it returns will mess up the rgb and effector_translation

{'observation': {'effector_target_translation': TensorSpec(shape=(1, 1, 4, 4, 2), dtype=tf.float32, name=None),
  'rgb': TensorSpec(shape=(1, 1, 4, 4, 2), dtype=tf.float32, name=None),
  'instruction': TensorSpec(shape=(1, 1, 4, 4, 512), dtype=tf.int32, name=None),
  'effector_translation': TensorSpec(shape=(1, 1, 4, 4, 180, 320, 3), dtype=tf.float32, name=None),
  'instruction_tokenized_clip': TensorSpec(shape=(1, 1, 4, None, 77), dtype=tf.int64, name=None)},
 'action': TensorSpec(shape=(1, 1, 4, 2), dtype=tf.float32, name=None)}