Implement a tf.data.Dataset for handling the input observations

[tobias-liaudat]

This will reduce considerably the memory footprint of the model as it won't require to load the entire dataset in GPU memory. Consequently, it will allow us to train the model in smaller GPUs and/or using more stars.