[laughingman2]

Have you checked RSSM approach in DreamerV1,V2,V3,PlaNet? It uses deterministic (GRU hidden state) and discrete stochastic latent states. The deterministic and stochastic (sampled) latent state are used to predict the next state. I think the stochastic state might help with your problem a bit.

[sillysaurusx]

Dear mystery HN’er, thank you so much. I hadn’t heard about RSSM, and your explanation was wonderfully helpful.

Much appreciated. Have a great weekend :)