Variational auto-encoder for learning dynamics model

This VAE network can predict the next states of the environment given the current state and action. Tested on ai safety grid world environments by DeepMind. I have written a gym wrapper for the same here Gym AI Safety Gridworlds

Results

After 50 epochs

States

True next states

Predicted next states

After 1000 epochs

States

True next states

Predicted next states

After 1800 epochs

States

True next states

Predicted next states