[论文解读] Koopman-Assisted Reinforcement Learning

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

EEG and Brain-Computer InterfacesNeuroscience被引用 3

一句话总结

本文提出了一种新型框架——Koopman辅助强化学习（KARL），该框架利用Koopman算子将非线性动力系统提升至无限维函数空间中的线性表示形式，从而更易于求解Hamilton-Jacobi-Bellman方程。通过使用Koopman张量重新表述软值迭代与软演员-critic算法，KARL在非线性、混沌及随机控制任务中实现了最先进性能，优于标准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.

研究动机与目标

• 为解决强化学习中高维非线性系统常见的Hamilton-Jacobi-Bellman方程求解困难问题。
• 通过将非线性动力学嵌入线性Koopman框架，提升深度强化学习的可解释性与样本效率。
• 将最大熵强化学习算法（如软Q学习与SAC）扩展至Koopman嵌入空间中运行。
• 实现对具有挑战性的系统（包括混沌与随机动力系统）的鲁棒控制，提升泛化能力与性能表现。

提出的方法

• 将非线性动力系统提升至Koopman不变函数空间，使系统动力学近似线性化。
• 将值函数表示为Koopman可观测量，从而通过Koopman算子实现值估计的线性传播。
• 构建一个编码状态与控制输入联合依赖关系的Koopman张量，以参数化贝尔曼更新。
• 利用Koopman张量重新表述软值迭代与软演员-critic算法，确保在提升空间中保持最大熵目标。
• 通过动态模态分解或扩展DMD等数据驱动方法，从轨迹中估计Koopman算子。
• 在软Koopman值迭代中采用离散化动作空间，以实现期望值更新的可计算性。

实验结果

研究问题

• RQ1 Koopman算子能否用于线性化非线性马尔可夫决策过程中的值函数动力学，从而实现更高效、更稳定的强化学习？
• RQ2 如何将Koopman算子整合进最大熵强化学习算法（如软Q学习与SAC）中，以在提升样本效率的同时保持探索能力？
• RQ3 基于Koopman的方法在非线性、混沌及随机系统上，相较于标准深度强化学习与经典控制基线方法，性能提升程度如何？
• RQ4 Koopman框架在复杂动力系统中如何增强对所学策略与值函数的可解释性？

主要发现

• KARL在四个基准控制任务上实现了最先进性能：线性状态空间系统、Lorenz系统、绕圆柱体流动的流体、以及具有非各向同性随机扰动的双阱势能系统。
• 在非线性与混沌系统上，KARL优于基于神经网络的标准SAC与经典LQR基线方法。
• Koopman辅助的软值迭代与演员-critic算法在学习最优策略时展现出更优的样本效率与稳定性。
• 消融实验表明，Koopman张量公式是性能的关键，若移除Koopman提升，性能将显著下降。
• 该框架通过揭示主导可观测量（特征），显著提升了高维系统中值函数动力学的可解释性。
• 对连续动作的初步结果表明，通过高斯空间中的KL近似软策略，可实现连续动作的扩展，完整实现留待未来工作。