Action reconstruction of skill vae

[OrangeSodahub]

Hi, thanks for your excellent opensource work! I run the training of the skill_vae, and visualize the reconstructed actions and original actions:

where the red ones are the outputs of skill_vae. However, it looks like these two actions are not perfectly overlapped. I would like to know if that is correct? Thanks!

[OrangeSodahub]

And the 2) question is, when I use smaller skill_block_size, e.g., from 32 to 16, while keep downsample factor to be unchanged 4, is it good and what changes of fsq should I make? I think I should have smaller codebook size? and latent dimension?

[OrangeSodahub]

Hi, @atharvamete Could you check this? Thanks!

[atharvamete]

Hi @OrangeSodahub,
Apologies for getting back late. What was the final reconstruction loss at the end of VAE training? Had it been converged? Despite convergence depending on dataset/action space/normalization the reconstruction error can vary. Any VAE will always have some reconstruction error so I would not expect trajectories to perfectly align, though lower the better.
Some tips to improve reconstruction loss:

1. use a downsampling factor of 2 instead of 4 (atleast on libero, this will improve reconstruction loss but might give lower success rate)
2. use skill_block_size of approx 1 to 1.5 second of motion as a thumb rule, (though on libero chunk size as large as 64 did not affect the performance much)
3. you can try increasing the vocabulory size to 2k or 4k depending on your dataset. On LIBERO larger ones gave lower recon loss but slightly lower success rate as well. The latent dimension is determined by the vocabulory size in case of FSQ.

I believe training the VAE does/should not take long since only action data is involved so feel free to play around with these parameters.

Lastly, reconstruction loss is not a reliable indicator of actual success rate, it's enough to ensure that it's low enough (<0.05 L1 loss on libero). Lemme know if you have more questions! All the best.

[OrangeSodahub]

@atharvamete Thanks, however, I was wondering if you did evaluate the reconstructed actions of skill_vae? The results above are from the checkpoint I trained on your codebase, following your original settings: block_size 32, downsample factor 4, train on libero90.

Have you checked that? Thanks!

Additionally, I've adjusted some other parameters, however I get all similar outputs.

[OrangeSodahub]

@atharvamete Or could you release your checkpoint like skill_vae part? I just want to check the reconstruction performance of it.

[atharvamete]

Oh I see, if that's the reconstruction you get on the default parameters then that should be what we might have had, does this affect success rate though? I found a checkpoint you can try.

[Kalcy-U]

Same issue. The reconstructed action trajectory basically cannot be successfully replayed. Is this normal?