Probabilistic Dreaming for World Models

“Dreaming” enables agents to learn from imagined experiences, enabling more robust and sample-efficient learning of world models. In this work, we consider innovations to the state-of-the-art Dreamer model using probabilistic methods that enable: (1) the parallel exploration of many latent states; and (2) maintaining distinct hypotheses for mutually exclusive futures while retaining the desirable gradient properties of continuous latents. Evaluating on the MPE SimpleTag domain, our method outperforms standard Dreamer with a 4.5% score improvement and 28% lower variance in episode returns. We also discuss limitations and directions for future work, including how optimal hyperparameters (e.g. particle count K) scale with environmental complexity, and methods to capture epistemic uncertainty in world models.

💡 Research Summary

The paper introduces “Probabilistic Dreaming” as an extension of the Dreamer‑v3 model‑based reinforcement learning framework. Traditional Dreamer learns a latent dynamics model (RSSM) with continuous Gaussian latents but, during imagination, samples only a single latent state to generate one imagined trajectory. This single‑sample approach limits exploration of multimodal futures and can collapse distinct possibilities into an unrealistic average.

To address these issues, the authors propose three key innovations:

1. Particle Filter for Latent Distribution – Instead of a single Gaussian sample, a set of K particles ({h_k, z_k}_{k=1}^K) is maintained during imagination. Each particle follows the Gaussian transition defined by the prior, and after stochastic propagation and resampling the empirical particle set approximates a potentially multimodal belief over latent states.

2. Latent Beam Search – For each particle, N candidate actions are sampled from the policy (\pi_\theta(a|h_k, z_k)). This creates (K \times N) branches that are simultaneously rolled out through the world model, allowing parallel exploration of many action‑state trajectories.

3. Free‑Energy‑Based Pruning – Because imagined trajectories lack real observations, the authors cannot use maximum‑likelihood to prune particles. Instead they score each branch with a free‑energy objective:
   \