Question about KL collapse in LAM training (beta=1e-6) — is the model effectively deterministic?

[soul667]

Hi, thanks for open-sourcing this project.

I have a question regarding the latent action model (LAM) training objective and the role of the KL term.

From the config:

beta: 0.000001

and from the implementation in lam/modules/lam.py:

if not self.training:
    z_rep = z_mu
else:
    z_rep = z_mu + torch.randn_like(z_var) * torch.exp(0.5 * z_var)

it seems that:

Training formally uses VAE-style reparameterization
But since beta = 1e-6, the KL regularization is almost negligible
As a result, the model can minimize noise injection by driving z_logvar -> -inf

Therefore:

exp(0.5 * z_logvar) -> 0

and sampling effectively collapses to:

z_rep ≈ z_mu

So my understanding is that, although the implementation is written as a VAE, the actual training dynamics are closer to a deterministic autoencoder with a 32-dimensional bottleneck.

This also seems consistent with inference using:

z_rep ≈ z_mu

My questions are:

Is this interpretation correct?
Was the extremely small beta intentionally chosen to encourage deterministic latent actions and maximize reconstruction quality?
In practice, does the KL term contribute anything meaningful during training?
Did you experiment with larger beta values or a true stochastic latent space?
Is the goal here representation learning rather than generative latent action modeling?

Thanks!