# Uncertainty Representations in State-Space Layers for Deep Reinforcement Learning under Partial Observability

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 52 of its 81 references that are in the map.
- First posted on arXiv 25 Sep 2024.

## Placed in

- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [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 Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Soft Actor-Critic Algorithms and Applications](https://mlsanity.com/paper/1812.05905/) (2018)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [HiPPO: Recurrent Memory with Optimal Polynomial Projections](https://mlsanity.com/paper/2008.07669/) (2020)
- [Hungry Hungry Hippos: Towards Language Modeling with State Space Models](https://mlsanity.com/paper/2212.14052/) (2022)
- [Deep Recurrent Q-Learning for Partially Observable MDPs](https://mlsanity.com/paper/1507.06527/) (2015)
- [On the Parameterization and Initialization of Diagonal State Space Models](https://mlsanity.com/paper/2206.11893/) (2022)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Target-driven Visual Navigation in Indoor Scenes using Deep Reinforcement Learning](https://mlsanity.com/paper/1609.05143/) (2016)
- [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)
- [Mega: Moving Average Equipped Gated Attention](https://mlsanity.com/paper/2209.10655/) (2022)
- [Soft Actor-Critic for Discrete Action Settings](https://mlsanity.com/paper/1910.07207/) (2019)
- [It's Raw! Audio Generation with State-Space Models](https://mlsanity.com/paper/2202.09729/) (2022)
- [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)

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

- [Kalman Linear Attention: Parallel Bayesian Filtering For Efficient Language Modelling and State Tracking](https://mlsanity.com/paper/2602.10743/) (2026)
- [Learning State-Tracking from Code Using Linear RNNs](https://mlsanity.com/paper/2602.14814/) (2026)
- [Long-Horizon Model-Based Offline Reinforcement Learning Without Explicit Conservatism](https://mlsanity.com/paper/2512.04341/) (2025)
- [ROSA-RL: Uncertainty-Aware Roundabout Optimized Speed Advisory with Reinforcement Learning](https://mlsanity.com/paper/2606.16558/) (2026)
- [Adversarial observations in probabilistic State-Space Models for robust Reinforcement Learning](https://mlsanity.com/paper/2606.20880/) (2026)

## Abstract (arXiv)

Optimal decision-making under partial observability requires reasoning about the uncertainty of the environment's hidden state. However, most reinforcement learning architectures handle partial observability with sequence models that have no internal mechanism to incorporate uncertainty in their hidden state representation, such as recurrent neural networks, deterministic state-space models and transformers. Inspired by advances in probabilistic world models for reinforcement learning, we propose a standalone Kalman filter layer that performs closed-form Gaussian inference in linear state-space models and train it end-to-end within a model-free architecture to maximize returns. Similar to efficient linear recurrent layers, the Kalman filter layer processes sequential data using a parallel scan, which scales logarithmically with the sequence length. By design, Kalman filter layers are a drop-in replacement for other recurrent layers in standard model-free architectures, but importantly they include an explicit mechanism for probabilistic filtering of the latent state representation. Experiments in a wide variety of tasks with partial observability show that Kalman filter layers excel in problems where uncertainty reasoning is key for decision-making, outperforming other stateful models.