Unable to generate decent images with VAE

Hey there,

I want to replicate some of the results obtained in the World Model paper. I created a dataset of the gym environment RacingCar and I trained a VAE model on it. However, the loss function seems stuck and it is unable to sample nice images. What bothers me the most is that all images are very similar to each other.

I checked the sanity of the dataset, but it seems okay. I believe my implementation (https://github.com/Mehd6384/Generative-Networks/blob/master/CarRacingVAE.py) is really close to the other ones out there. Truth is, I can’t figure out what’s wrong ? worstbest

(Above, the only two kind of images I get)

Has anyone been stuck on VAE with similar results ?

Thanks a lot!