Imagination-Augmented in practice

Question

I am doing a research on Imagination-Augmented, would you like to understand in practice how this using Tensorflow? I would like to know the recommended functions to work in this way, and the logic itself.