[Paper Review] Koopman-Assisted Reinforcement Learning

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

EEG and Brain-Computer InterfacesNeuroscience3 citations

TL;DR

This paper introduces Koopman-Assisted Reinforcement Learning (KARL), a novel framework that leverages the Koopman operator to lift nonlinear dynamical systems into a linear representation in an infinite-dimensional function space, enabling more tractable solution of the Hamilton-Jacobi-Bellman equation. By reformulating soft value iteration and soft actor-critic using a Koopman tensor, KARL achieves state-of-the-art performance on nonlinear, chaotic, and stochastic control tasks, outperforming both standard SAC and classical LQR baselines.

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.

Motivation & Objective

• To address the intractability of solving the Hamilton-Jacobi-Bellman equation in high-dimensional, nonlinear systems common in reinforcement learning.
• To improve interpretability and sample efficiency in deep RL by embedding nonlinear dynamics into a linear Koopman framework.
• To extend maximum-entropy RL algorithms like soft Q-learning and SAC to operate in Koopman-embedded spaces.
• To enable robust control of challenging systems, including chaotic and stochastic dynamical systems, with improved generalization and performance.

Proposed method

• Lift nonlinear dynamics into a Koopman-invariant function space where the system dynamics become approximately linear.
• Represent the value function as a Koopman observable, enabling linear propagation of value estimates via the Koopman operator.
• Construct a Koopman tensor that encodes the joint dependence of dynamics on state and control inputs, enabling parameterization of the Bellman update.
• Reformulate soft value iteration and soft actor-critic using the Koopman tensor to maintain maximum entropy objectives in the lifted space.
• Use data-driven Koopman operator approximation via dynamic mode decomposition or extended DMD to estimate the Koopman operator from trajectories.
• Implement a discretized action space in the soft Koopman value iteration for tractable computation of expected value updates.

Experimental results

Research questions

• RQ1 Can the Koopman operator be used to linearize the value function dynamics in nonlinear Markov Decision Processes, enabling more efficient and stable RL?
• RQ2 How can the Koopman operator be integrated into maximum-entropy RL algorithms like soft Q-learning and SAC to maintain exploration while improving sample efficiency?
• RQ3 To what extent can Koopman-based methods outperform standard deep RL and classical control baselines on nonlinear, chaotic, and stochastic systems?
• RQ4 How does the Koopman framework enhance interpretability of learned policies and value functions in complex dynamical systems?

Key findings

• KARL achieves state-of-the-art performance on four benchmark control tasks: a linear state-space system, the Lorenz system, fluid flow past a cylinder, and a double-well potential with non-isotropic stochastic forcing.
• KARL outperforms standard neural network-based Soft Actor-Critic (SAC) and classical Linear Quadratic Regulator (LQR) baselines on nonlinear and chaotic systems.
• The Koopman-assisted soft value iteration and actor-critic algorithms show superior sample efficiency and stability in learning optimal policies.
• Ablation studies confirm that the Koopman tensor formulation is essential for performance, with significant degradation when the Koopman lift is removed.
• The framework enables improved interpretability by revealing the dominant observables (features) that drive value function dynamics, especially in high-dimensional systems.
• Preliminary results on continuous actions suggest that KL-approximation of the soft policy in Gaussian space can enable continuous action extensions, with full implementation left for future work.