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

TL;DR

KalMamba introduces an efficient probabilistic state space model architecture for reinforcement learning that combines the scalability of deterministic models with the benefits of probabilistic inference, enabling faster processing of long sequences.

Contribution

KalMamba is a novel architecture that integrates probabilistic SSMs with deterministic scalable methods, improving efficiency in RL tasks involving high-dimensional, partial data.

Findings

KalMamba achieves competitive performance with state-of-the-art SSMs in RL.

It significantly enhances computational efficiency, especially on long sequences.

The model leverages parallel associative scanning for scalable inference.

Abstract

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.