KalMamba: Towards Efficient Probabilistic State Space Models for RL under Uncertainty
Fuente:
arXiv
Saved in:
| Main Authors: | , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866929394237308928 |
|---|---|
| author | Becker, Philipp Freymuth, Niklas Neumann, Gerhard |
| author_facet | Becker, Philipp Freymuth, Niklas Neumann, Gerhard |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2406_15131 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | KalMamba: Towards Efficient Probabilistic State Space Models for RL under Uncertainty Becker, Philipp Freymuth, Niklas Neumann, Gerhard Machine Learning 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. |
| title | KalMamba: Towards Efficient Probabilistic State Space Models for RL under Uncertainty |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2406.15131 |