Data loader normalization fn issue

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)}
```