Autoencoder vs VAE [question]

[el-cangrejo]

Hello, I trained two state representation models, one with the vae loss and one with the autoencoder loss, both on the KukaRandButtonGymEnv-v0 and for the same number of epochs (5) and state dims (5). The vae model reconstructs very poorly the input image compared to the autoencoder:

vae:

autoencoder:

It seems like it's producing an average over the dataset. Does it sample from the data likelihood distribution or does it calculate the mean value of p(x|z)?

Also when visualizing the learned space with the enjoy_latent module I noticed that the states of the autoencoder have a wide range (for example [-20, 20]) and the produced visualizations are quite different. On the other hand the states for the vae have a very narrow range (for example [-0.04, 0.04]) and the produced visualizations are almost identical. Is this an expected behavior?

Autoencoder (min, max):

VAE (min, max):

[araffin]

Hello,

The vae model reconstructs very poorly the input image compared to the autoencoder:
Also when visualizing the learned space with the enjoy_latent module I noticed that the states of the autoencoder have a wide range

I would say those two behaviors are expected. The VAE has an additional KL loss, so its latent space distribution approximates a multi-dimensional gaussian with a unit standard deviation. This loss is combined with the reconstruction loss, so they may interfere: as it adds a new constrain, it is harder for the VAE to reconstruct the image. The KL loss also explains the bounds of the VAE latent space: the AE is not penalized when its latent space has a wide range, but this is true for the VAE.

In your case, I would try to add more dimensions (so it is easier for the VAE to reconstruct the image), more epochs, and maybe play with the \beta coefficient (which weights the importance of the KL loss and reconstruction loss). Btw, how many samples are you using for training?

Does it sample from the data likelihood distribution or does it calculate the mean value of p(x|z)?

It samples, the corresponding line: https://github.com/araffin/srl-zoo/blob/master/evaluation/enjoy_latent.py#L35

[el-cangrejo]

Hi, thanks for the quick response. I'm using 1568 samples. I'll try with more dimensions, samples and epochs and see how it goes. Do you provide any way to visualize a "combined" model? I'm trying to implement something like the World Models model so I trained a model with the forward and autoencoder loss and I would like the the forward model to predict the next state given the current and an action and the autoencoder to reconstruct the observation from the state.

[araffin]

Do you provide any way to visualize a "combined" model?

We provide way to visualize combined models when combined means combining a method with a decoder (=auto-encoder) with another loss (ex: inverse dynamics, where we use knn on the dataset).
However, we don't provide visualization for future observation, but it should not be too hard to integrate.

Feel free to submit a PR if you integrate it ;)

[el-cangrejo]

Okay thanks!

[ncble]

Actually, about the VAE and autoencoder, I have recently discovered a bug (see the issue #53). By fixing this bug, I have successfully reconstructed image with VAE:

The left two images are original images, and the right two are reconstructed. Without my fix as described in #53, I got similar result to your images. So it's definitely the source of your problem.