# On Uncertainty in Deep State Space Models for Model-Based Reinforcement Learning

arXiv: https://arxiv.org/abs/2210.09256 · HTML: https://mlsanity.com/paper/2210.09256/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 53 references that are in the map.
- First posted on arXiv 17 Oct 2022.

## Builds on (most cited first, 30 of 39)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [On the Properties of Neural Machine Translation: Encoder-Decoder Approaches](https://mlsanity.com/paper/1409.1259/) (2014)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods](https://mlsanity.com/paper/1910.09457/) (2019)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Planning to Explore via Self-Supervised World Models](https://mlsanity.com/paper/2005.05960/) (2020)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Structured Inference Networks for Nonlinear State Space Models](https://mlsanity.com/paper/1609.09869/) (2016)
- [Deep Kalman Filters](https://mlsanity.com/paper/1511.05121/) (2015)
- [Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Data](https://mlsanity.com/paper/1605.06432/) (2016)
- [Algorithmic Framework for Model-based Deep Reinforcement Learning with Theoretical Guarantees](https://mlsanity.com/paper/1807.03858/) (2018)
- [Disentangled Sequential Autoencoder](https://mlsanity.com/paper/1803.02991/) (2018)
- [A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning](https://mlsanity.com/paper/1710.05741/) (2017)
- [Objective Mismatch in Model-based Reinforcement Learning](https://mlsanity.com/paper/2002.04523/) (2020)
- [Variational Sequential Monte Carlo](https://mlsanity.com/paper/1705.11140/) (2017)
- [Recurrent Kalman Networks: Factorized Inference in High-Dimensional Deep Feature Spaces](https://mlsanity.com/paper/1905.07357/) (2019)
- [Probabilistic Recurrent State-Space Models](https://mlsanity.com/paper/1801.10395/) (2018)
- [Learning and Querying Fast Generative Models for Reinforcement Learning](https://mlsanity.com/paper/1802.03006/) (2018)
- [Backprop KF: Learning Discriminative Deterministic State Estimators](https://mlsanity.com/paper/1605.07148/) (2016)
- [From Pixels to Torques: Policy Learning with Deep Dynamical Models](https://mlsanity.com/paper/1502.02251/) (2015)
- [Robust Locally-Linear Controllable Embedding](https://mlsanity.com/paper/1710.05373/) (2017)

## Built on by (most cited first, 13 of 13)

- [MUVO: A Multimodal Generative World Model for Autonomous Driving with Geometric Representations](https://mlsanity.com/paper/2311.11762/) (2023)
- [R-AIF: Solving Sparse-Reward Robotic Tasks from Pixels with Active Inference and World Models](https://mlsanity.com/paper/2409.14216/) (2024)
- [On Rollouts in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2501.16918/) (2025)
- [Combining Reconstruction and Contrastive Methods for Multimodal Representations in RL](https://mlsanity.com/paper/2302.05342/) (2023)
- [KalMamba: Towards Efficient Probabilistic State Space Models for RL under Uncertainty](https://mlsanity.com/paper/2406.15131/) (2024)
- [Exploring the Potential of World Models for Anomaly Detection in Autonomous Driving](https://mlsanity.com/paper/2308.05701/) (2023)
- [Uncertainty Representations in State-Space Layers for Deep Reinforcement Learning under Partial Observability](https://mlsanity.com/paper/2409.16824/) (2024)
- [Biased Dreams: Limitations to Epistemic Uncertainty Quantification in Latent Dynamics Models](https://mlsanity.com/paper/2604.25416/) (2026)
- [Learning World Models With Hierarchical Temporal Abstractions: A Probabilistic Perspective](https://mlsanity.com/paper/2404.16078/) (2024)
- [Adaptive World Models: Learning Behaviors by Latent Imagination Under Non-Stationarity](https://mlsanity.com/paper/2411.01342/) (2024)
- [All Models are Wrong, Knowing Where is Useful: On Model Uncertainty in Reinforcement Learning](https://mlsanity.com/paper/2606.01363/) (2026)
- [Reinformed Dreamer: An Asymmetric World Model Efficiently Trained through Latent Guidance](https://mlsanity.com/paper/2607.26040/) (2026)
- [Admissable: Training Reinforcement Learning Agents against Adversarial Missingness](https://mlsanity.com/paper/2609.15297/) (2026)

## Abstract (arXiv)

Improved state space models, such as Recurrent State Space Models (RSSMs), are a key factor behind recent advances in model-based reinforcement learning (RL). Yet, despite their empirical success, many of the underlying design choices are not well understood. We show that RSSMs use a suboptimal inference scheme and that models trained using this inference overestimate the aleatoric uncertainty of the ground truth system. We find this overestimation implicitly regularizes RSSMs and allows them to succeed in model-based RL. We postulate that this implicit regularization fulfills the same functionality as explicitly modeling epistemic uncertainty, which is crucial for many other model-based RL approaches. Yet, overestimating aleatoric uncertainty can also impair performance in cases where accurately estimating it matters, e.g., when we have to deal with occlusions, missing observations, or fusing sensor modalities at different frequencies. Moreover, the implicit regularization is a side-effect of the inference scheme and not the result of a rigorous, principled formulation, which renders analyzing or improving RSSMs difficult. Thus, we propose an alternative approach building on well-understood components for modeling aleatoric and epistemic uncertainty, dubbed Variational Recurrent Kalman Network (VRKN). This approach uses Kalman updates for exact smoothing inference in a latent space and Monte Carlo Dropout to model epistemic uncertainty. Due to the Kalman updates, the VRKN can naturally handle missing observations or sensor fusion problems with varying numbers of observations per time step. Our experiments show that using the VRKN instead of the RSSM improves performance in tasks where appropriately capturing aleatoric uncertainty is crucial while matching it in the deterministic standard benchmarks.