KalMamba: Towards Efficient Probabilistic State Space Models for RL under Uncertainty
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.
Code (0)
등록된 구현이 없습니다.
Tasks
Computational EfficiencyMambaReinforcement Learning (RL)State Space ModelsSimilar Papers 제목 키워드 기반
Uncertainty Representations in State-Space Layers for Deep Reinforcement Learning under Partial Observability
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 sequenc…
Decision MakingDeep Reinforcement Learningreinforcement-learningReinforcement Learning+1Bridging Learned Visual Perception and Symbolic Belief-Space Planning
In partially observable settings, agents must act without full knowledge of the world state and rely on uncertain state-estimation pipelines. Obtaining grounded and verifiable symbolic plans under such uncertainty remain…
Adversarial observations in probabilistic State-Space Models for robust Reinforcement Learning
Decision-making under partial or adversarial observability requires accurate inference of the environment's latent state and its associated uncertainty. This work analyses adversarial attacks on linear probabilistic stat…
Reinforcement LearningLearning Interpretable Deep State Space Model for Probabilistic Time Series Forecasting
Probabilistic time series forecasting involves estimating the distribution of future based on its history, which is essential for risk management in downstream decision-making. We propose a deep state space model for pro…
Decision MakingManagementProbabilistic Time Series ForecastingTime Series+2Unified Probabilistic Neural Architecture and Weight Ensembling Improves Model Robustness
Robust machine learning models with accurately calibrated uncertainties are crucial for safety-critical applications. Probabilistic machine learning and especially the Bayesian formalism provide a systematic framework to…
Bayesian InferenceNeural Architecture Search