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

要約

確率的状態空間モデル (SSM) は、制御のための簡潔な表現を提供するため、高次元の部分情報からの強化学習 (RL) に不可欠です。
しかし、S4 や Mamba などの最近の決定論的な対応物に比べて計算効率が劣ります。
我々は、確率的 SSM の強みと決定的 SSM のスケーラビリティを組み合わせた、RL の表現を学習するための効率的なアーキテクチャである KalMamba を提案します。
KalMamba は Mamba を活用して、潜在空間内の線形ガウス SSM のダイナミクス パラメーターを学習します。
この潜在空間での推論は、標準的なカルマン フィルタリングと平滑化に相当します。
私たちは、Mamba と同様の並列連想スキャンを使用してこれらの操作を実現し、原理に基づいた、高効率でスケーラブルな確率的 SSM を取得します。
私たちの実験は、KalMamba が RL における最先端の SSM アプローチと競合しながら、特に長い相互作用シーケンスでの計算効率を大幅に向上させることを示しています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google