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. On DeepMind Control Suite environments, the method obtains final returns comparable to DreamerV3, with faster learning on Hopper and slower learning on Walker. On Meta-World door-open, with goals resampled every episode, DreamerV3 records 0 successes in $5\times10^{5}$ steps, whereas our model solves the task with a return of 88\% of the scripted-expert return. The same inference procedure recovers the reference regimes with low error on two of four offline switching benchmarks, including one on which a tree-structured recurrent SLDS has substantially higher segmentation error. Our online results support a contribution from the sticky linear parameterisation, while the reported attempts to prevent collapse to one occupied regime do not restore multiple occupied regimes.
Chat is not available.
Successful Page Load