[論文レビュー] Koopman-Assisted Reinforcement Learning

Preston Rozwood, Edward Mehrez+3 more|arXiv (Cornell University)|Mar 4, 2024

EEG and Brain-Computer InterfacesNeuroscience被引用数 3

ひとこと要約

本稿では、非線形力学系を無限次元関数空間における線形表現に高次元化するためのコプマン作用素を活用する、Koopman-Assisted Reinforcement Learning (KARL) と呼ばれる新規フレームワークを提案する。これにより、ハミルトニアン・ジャコビ・ベルマン方程式のより取り扱いやすい解法が可能となる。Koopmanテンソルを用いてソフト値反復およびソフトアクトアクリティックを再定式化することで、非線形的でカオス的かつ確率的な制御タスクにおいて、標準的なSACおよび古典的LQRベースラインを上回る最先端の性能を達成する。

ABSTRACT

The Bellman equation and its continuous form, the Hamilton-Jacobi-Bellman (HJB) equation, are ubiquitous in reinforcement learning (RL) and control theory. However, these equations quickly become intractable for systems with high-dimensional states and nonlinearity. This paper explores the connection between the data-driven Koopman operator and Markov Decision Processes (MDPs), resulting in the development of two new RL algorithms to address these limitations. We leverage Koopman operator techniques to lift a nonlinear system into new coordinates where the dynamics become approximately linear, and where HJB-based methods are more tractable. In particular, the Koopman operator is able to capture the expectation of the time evolution of the value function of a given system via linear dynamics in the lifted coordinates. By parameterizing the Koopman operator with the control actions, we construct a ``Koopman tensor'' that facilitates the estimation of the optimal value function. Then, a transformation of Bellman's framework in terms of the Koopman tensor enables us to reformulate two max-entropy RL algorithms: soft value iteration and soft actor-critic (SAC). This highly flexible framework can be used for deterministic or stochastic systems as well as for discrete or continuous-time dynamics. Finally, we show that these Koopman Assisted Reinforcement Learning (KARL) algorithms attain state-of-the-art (SOTA) performance with respect to traditional neural network-based SAC and linear quadratic regulator (LQR) baselines on four controlled dynamical systems: a linear state-space system, the Lorenz system, fluid flow past a cylinder, and a double-well potential with non-isotropic stochastic forcing.

研究の動機と目的

• 強化学習において一般的に見られる高次元で非線形なシステムにおけるハミルトニアン・ジャコビ・ベルマン方程式の解法の非可解性に対処すること。
• 非線形力学を線形コプマンフレームワークに埋め込むことで、深層強化学習における解釈可能性とサンプル効率を向上させること。
• ソフトQ学習やSACのような最大エントロピー強化学習アルゴリズムを、コプマン埋め込み空間で動作可能にするように拡張すること。
• カオス的および確率的力学系を含む挑戦的なシステムに対するロバストな制御を可能にし、一般化性能と性能を向上させること。

提案手法

• 非線形力学を、システムの力学が近似的に線形となるコプマン不変関数空間へと高次元化する。
• 値関数をコプマン観測量として表現することで、コプマン作用素による値推定の線形的伝播を可能にする。
• 状態と制御入力の両者に依存する動的関係を符号化するコプマンテンソルを構築し、ベルマン更新のパラメータ化を可能にする。
• コプマンテンソルを用いてソフト値反復およびソフトアクトアクリティックを再定式化し、高次元空間における最大エントロピー目的を維持する。
• トレーブルからトレースを用いた動的モード分解または拡張型DMDによるデータ駆動型コプマン作用素の近似を用いて、コプマン作用素を推定する。
• 計算の tractability を確保するため、ソフトコプマン値反復において離散化された行動空間を実装する。

実験結果

リサーチクエスチョン

• RQ1 コプマン作用素を用いることで、非線形マルコフ決定過程における値関数ダイナミクスを線形化し、より効率的かつ安定した強化学習を可能にすることができるか？
• RQ2 コプマン作用素を最大エントロピー強化学習アルゴリズム（例：ソフトQ学習やSAC）に統合することで、探索を維持しつつサンプル効率を向上させることができるか？
• RQ3 コプマンベースの手法は、非線形的でカオス的かつ確率的なシステムにおいて、標準的な深層強化学習および古典的制御ベースラインをどの程度上回ることができるか？
• RQ4 コプマンフレームワークは、複雑な力学系における学習済みポリシーおよび値関数の解釈可能性をどの程度向上させるか？

主な発見

• KARLは、線形状態空間系、ローレンツ系、円柱周囲の流れ、非等方的確率的力が作用するダブルウェルポテンシャルの4つのベンチマーク制御タスクで最先端の性能を達成した。
• 非線形的およびカオス的システムにおいて、標準的なニューラルネットワークベースのソフトアクトアクリティック（SAC）および古典的線形二次調節器（LQR）ベースラインを上回った。
• コプマン補助ソフト値反復およびアクトアクリティックアルゴリズムは、最適ポリシーの学習において優れたサンプル効率と安定性を示した。
• アブレーションスタディの結果、コプマンテンソルの定式化が性能に不可欠であることが確認され、コプマン高次元化を除去すると著しい性能低下が生じた。
• 特に高次元系において、価値関数ダイナミクスを駆動する主要な観測量（特徴）を明らかにすることで、解釈可能性の向上が達成された。
• 連続的行動への応用に関する予備的結果では、ガウス空間におけるソフトポリシーのKL近似が連続的行動拡張を可能にする可能性を示唆しており、完全な実装は今後の課題である。