Dream to Explore: Discovering Latent Dynamical Regimes with Sticky Nonparametric World Models
Zahra Sheikhbahaee ⋅ Gopeshh Subbaraj ⋅ Adam Safron ⋅ Saurav Jha ⋅ Irina Rish ⋅ Sarath Chandar
Abstract
World-model agents predict future latent states using an amortized neural transition prior conditioned on a recurrent state. In DreamerV3, this prior is expressive and can represent different dynamical contexts implicitly through the recurrent state. However, it does not explicitly partition experience into persistent regimes whose dynamics can be adapted separately as the data distribution changes. We replace DreamerV3's discrete latent and neural transition prior with a continuous latent whose prior is a Bayesian nonparametric mixture of linear regimes. Each regime has its own action response, and a sticky transition rule promotes persistence across time steps. We fit the regime model using structured variational Bayes with global sufficient statistics blended at a constant rate, and we derive the corresponding exact recency-weighted conjugate update for fixed variational messages. The model supports additions and removals of regimes during training to adapt its structure. We compare it with DreamerV3 trained in the same codebase, R2I and PPO. On Meta-World manipulation with goals resampled every episode, our model has the highest mean return after $5\times10^{5}$ environment steps on reach, door-open and drawer-open, from both proprioceptive state and pixels; on door-open from pixels it reaches about $97\%$ of the scripted-expert return, compared with $51\%$ for DreamerV3. Results on DeepMind Control are task dependent: the method learns faster on Hopper Hop but underperforms DreamerV3 on a few other tasks. Offline, it has the lowest mean segmentation error on NASCAR and substantially lower error than TrSLDS on ToyARK13, while rSLDS has a slightly lower mean error on the latter.
Chat is not available.
Successful Page Load