Koopman-Assisted Reinforcement Learning

TL;DR

This paper introduces Koopman-assisted reinforcement learning algorithms that leverage data-driven Koopman operators to linearize complex systems, enabling more tractable and interpretable control solutions with state-of-the-art results.

Contribution

The paper develops novel RL algorithms based on the Koopman operator, allowing for linearization of nonlinear dynamics and improved performance over traditional neural network methods.

Findings

Achieves state-of-the-art performance on various nonlinear systems.

Successfully reformulates RL algorithms using Koopman-based linear dynamics.

Handles both deterministic and stochastic, as well as discrete and continuous systems.

Abstract

The Bellman equation and its continuous form, the Hamilton-Jacobi-Bellman equation, are ubiquitous in reinforcement learning and control theory. However, these equations become intractable for high-dimensional or nonlinear systems. This paper develops two new reinforcement learning algorithms based on the data-driven Koopman operator, which lifts a nonlinear system into new coordinates where the dynamics become approximately linear, and where Hamilton-Jacobi-Bellman-based methods are more tractable. In particular, the Koopman operator captures the expectation of the time evolution of the value function via linear dynamics in the lifted coordinates. By parameterizing the Koopman operator with the control actions, we construct a ``controlled Koopman tensor'' that facilitates the estimation of the optimal value function. This enables us to reformulate two max-entropy RL algorithms: soft value iteration and soft actor-critic. This flexible and interpretable framework includes deterministic and stochastic systems, as well as discrete and continuous dynamics. Koopman Assisted reinforcement learning attains state-of-the-art performance with respect to traditional neural network-based soft actor-critic baselines on a linear state-space system, the Lorenz system, fluid flow past a cylinder, and a double-well potential with non-isotropic stochastic forcing.