mlsanityData to 7 Oct 2026

KalMamba: Towards Efficient Probabilistic State Space Models for RL under Uncertainty

arXiv 2406.15131 ↗First posted 21 Jun 2024cs.LG stat.MLcited by 9

Not on the map. Why?

Abstract (arXiv)

Probabilistic State Space Models (SSMs) are essential for Reinforcement Learning (RL) from high-dimensional, partial information as they provide concise representations for control. Yet, they lack the computational efficiency of their recent deterministic counterparts such as S4 or Mamba. We propose KalMamba, an efficient architecture to learn representations for RL that combines the strengths of probabilistic SSMs with the scalability of deterministic SSMs. KalMamba leverages Mamba to learn the dynamics parameters of a linear Gaussian SSM in a latent space. Inference in this latent space amounts to standard Kalman filtering and smoothing. We realize these operations using parallel associative scanning, similar to Mamba, to obtain a principled, highly efficient, and scalable probabilistic SSM. Our experiments show that KalMamba competes with state-of-the-art SSM approaches in RL while significantly improving computational efficiency, especially on longer interaction sequences.

Key facts. Cited by 9 papers in the map (data through 7 Oct 2026). Builds on 45 of its 55 references that are in the map. First posted on arXiv 21 Jun 2024.

Builds on

45 of 55 references.

References · 45

  1. 2017Attention Is All You Needcited by 59,932
  2. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  3. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
  4. 2015Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learningcited by 4,416
  5. 2013Auto-Encoding Variational Bayescited by 3,980
  6. 2021Efficiently Modeling Long Sequences with Structured State Spacescited by 2,424
  7. 2019Dream to Control: Learning Behaviors by Latent Imaginationcited by 1,598
  8. 2018World Modelscited by 1,444
  9. 2014On the Properties of Neural Machine Translation: Encoder-Decoder Approachescited by 1,303
  10. 2018Learning Latent Dynamics for Planning from Pixelscited by 1,221
  11. 2023Mastering Diverse Domains through World Modelscited by 1,193
  12. 2018Soft Actor-Critic Algorithms and Applicationscited by 1,179
33 more
  1. 2020Mastering Atari with Discrete World Modelscited by 990
  2. 2018DeepMind Control Suitecited by 935
  3. 2021Deep Reinforcement Learning at the Edge of the Statistical Precipicecited by 781
  4. 2022Simplified State Space Layers for Sequence Modelingcited by 776
  5. 2018Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Modelscited by 748
  6. 2022DayDreamer: World Models for Physical Robot Learningcited by 372
  7. 2022Diagonal State Spaces are as Effective as Structured State Spacescited by 350
  8. 2015Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Imagescited by 266
  9. 2019Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Modelcited by 228
  10. 2016Structured Inference Networks for Nonlinear State Space Modelscited by 182
  11. 2016Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Datacited by 154
  12. 2022TransDreamer: Reinforcement Learning with Transformer World Modelscited by 143
  13. 2022Liquid Structural State-Space Modelscited by 134
  14. 2018Disentangled Sequential Autoencodercited by 117
  15. 2017A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learningcited by 107
  16. 2017Variational Sequential Monte Carlocited by 93
  17. 2019Recurrent Kalman Networks: Factorized Inference in High-Dimensional Deep Feature Spacescited by 64
  18. 2021Temporal Predictive Coding For Model-Based Planning In Latent Spacecited by 54
  19. 2016Backprop KF: Learning Discriminative Deterministic State Estimatorscited by 47
  20. 2024Mastering Memory Tasks with World Modelscited by 38
  21. 2017Robust Locally-Linear Controllable Embeddingcited by 37
  22. 2022Deep Latent State Space Models for Time-Series Generationcited by 37
  23. 2015Neural Adaptive Sequential Monte Carlocited by 34
  24. 2015Black box variational inference for state space modelscited by 31
  25. 2017Identification of Gaussian Process State Space Modelscited by 31
  26. 2019Temporal Parallelization of Bayesian Smootherscited by 31
  27. 2021Learning Representations for Pixel-based Control: What Matters and Why?cited by 24
  28. 2019Switching Linear Dynamics for Variational Bayes Filteringcited by 17
  29. 2020Action-Conditional Recurrent Kalman Networks For Forward and Inverse Dynamics Learningcited by 13
  30. 2022On Uncertainty in Deep State Space Models for Model-Based Reinforcement Learningcited by 13
  31. 2022Hidden Parameter Recurrent State Space Models For Changing Dynamics Scenarioscited by 8
  32. 2018Deep State Space Models for Unconditional Word Generationcited by 6
  33. 2021Robust Robotic Control from Pixels using Contrastive Recurrent State-Space Modelscited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 2; 2025: 2; 2026: 5.

9 papers.

  1. 2025Advancing Intelligent Sequence Modeling: Evolution, Trade-offs, and Applications of State-Space Architectures from S4 to Mambacited by 9
  2. 2026Kalman Linear Attention: Parallel Bayesian Filtering For Efficient Language Modelling and State Trackingcited by 7
  3. 2024Uncertainty Representations in State-Space Layers for Deep Reinforcement Learning under Partial Observabilitycited by 5
  4. 2024Unraveling the Complexity of Memory in RL Agents: an Approach for Classification and Evaluationcited by 5
  5. 2025MEMBOT: Memory-Based Robot in Intermittent POMDPcited by 1
  6. 2026Geometry of Uncertainty: Learning Metric Spaces for Multimodal State Estimation in RLcited by 0
  7. 2026VariBASed: Variational Bayes-Adaptive Sequential Monte-Carlo Planning for Deep Reinforcement Learningcited by 0
  8. 2026Belief-State RWKV for Reinforcement Learning under Partial Observabilitycited by 0
  9. 2026How Wrong Can a Good Predictor Be? Diverging Updates with Vanishing Predictive KLcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2024Decision Mamba: A Multi-Grained State Space Model with Self-Evolution Regularization for Offline RL
  2. 2024Mamba as Decision Maker: Exploring Multi-scale Sequence Modeling in Offline Reinforcement Learning
  3. 2022On Practical Reinforcement Learning: Provable Robustness, Scalability, and Statistical Efficiency
  4. 2025Accelerating Model-Based Reinforcement Learning with State-Space World Models
  5. 2024Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling
  6. 2024Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient
  7. 2020Offline Reinforcement Learning from Images with Latent Space Models
  8. 2021Provably Efficient Representation Selection in Low-rank Markov Decision Processes: From Online to Offline RL