Koopman Q-learning: Offline Reinforcement Learning via Symmetries of Dynamics

1 INTRODUCTION . The recent impressive advances in reinforcement learning ( RL ) range from robotics , to strategy games and recommendation systems ( Kalashnikov et al. , 2018 ; Li et al. , 2010 ) . Reinforcement learning is canonically regarded as an active learning process - also referred to as online RL - where the agent interacts with the environment at each training run . In contrast , offline RL algorithms learn from large , previously collected static datasets , and thus do not rely on environment interactions ( Agarwal et al. , 2020a ; Ernst et al. , 2005 ; Fujimoto et al. , 2019 ) . Online data collection is performed by simulations or by means of real world interactions e.g . robotics and in either scenario interactions maybe costly and/or dangerous . In principle offline datasets only need to be collected once which alleviates the before-mentioned short-comings of costly online interactions . Offline datasets are typically collected using behavioral policies for the specific task ranging from , random policies , or near-optimal policies to human demonstrations . In particular , being able to leverage the latter is a major advantage of offline RL over online approaches , and then the learned policies can be deployed or finetuned on the desired environment . Offline RL has successfully been applied to learn agents that outperform the behavioral policy used to collect the data ( Kumar et al. , 2020 ; Wu et al. , 2019 ; Agarwal et al. , 2020b ; Ernst et al. , 2005 ) . However algorithms admit major shortcomings in regard to over-fitting and overestimating the true state-action values of the distribution . One solution was recently propsed by Sinha et al . ( 2021 ) , where they tested several data augmentation schemes to improve the performance and generalization capabilities of the learned policies . However , despite the recent progress , learning from offline demonstrations is a tedious endeavour as the dataset typically does not cover the full state-action space . Moreover , offline RL algorithms per definition do not admit the possibility for further environment exploration to refine their distributions towards an optimal policy . It was argued previously that is basically impossible for an offline RL agent to learn an optimal policy as the generalization to near data generically leads to compounding errors such as overestimation bias ( Kumar et al. , 2020 ) . In this paper , we look at offline RL through the lens of Koopman spectral theory in which nonlinear dynamics are represented in terms of a linear operator acting on the space of measurement functions of the system . Through which the representation the symmetries of the dynamics may be inferred directly , and can then be used to guide data augmentation strategies see Figure 1 . We further provide theoretical results on the existence on nature of symmetries relevant for control systems such as reinforcement learning . More specifically , we apply Koopman spectral theory by : first learning symmetries of the system ’ s underlying dynamic in a self-supervised fashion from the static dataset , and second employing the latter to extend the offline dataset at training time by out-of-distribution values . As this reflects the system ’ s dynamics the additional data is to be interpreted as an exploration of the environment ’ s phase space . Some prior works have explored symmetry of the state-action space in the context of Markov Decision Processes ( MDP ’ s ) ( Higgins et al. , 2018 ; Balaraman & Andrew , 2004 ; van der Pol et al. , 2020 ) since many control tasks exhibit apparent symmetries e.g . the classic cart-pole task which is symmetric across y-axis . However , the paradigm we introduce in this work is of a different nature entirely . The distinction is twofold : first , the symmetries are learned in a self-supervised way and are in general not apparent to the developer ; second : we concern with symmetry transformation from state tuples ( st , st+1 ) → ( s̃t , s̃t+1 ) which leave the action invariant inferred from the dynamics inherited by the behavioral policy of the underlying offline data . In other words we seek to derive a neighbourhood around a MDP tuple in the offline dataset in which the behavioral policy is likely to choose the same action based on its dynamics in the environment . In practice the Koopman latent space representation is learned in a self-supervised manner by training to predict the next state using a VAE model ( Kingma & Welling , 2013 ) . To summarize , in this paper , we propose Koopman Forward ( Conservative ) Q-learning ( KFC ) : a model-free Q-learning algorithm which uses the symmetries in the dynamics of the environment to guide data augmentation strategies . We also provide thorough theoretical justifications for KFC . Finally , we empirically test our approach on several challenging benchmark datasets from D4RL ( Fu et al. , 2021 ) , MetaWorld ( Yu et al. , 2019 ) and Robosuite ( Zhu et al. , 2020 ) and find that by using KFC we can improve the state-of-the-art on most benchmark offline reinforcement learning tasks . 2 PRELIMINARIES AND BACKGROUND . 2.1 OFFLINE RL & CONSERVATIVE Q-LEARNING . Reinforcement learning algorithms train policies to maximize the cumulative reward received by an agent who interacts with an environment . Formally the setting is given by a Markov decision process ( S , A , ρ , r , γ ) , with state space S , action space A , and ρ ( st+1|st , at ) the transition density function from the current state and action to the next state . Moreover , γ is the discount factor and r ( St ) the reward function . At any discrete time the agent chooses an action at ∈ A according to its underlying policy πθ ( at|st ) based on the information of the current state st ∈ S where the policy is parametrized by θ . We focus on the Actor-Critic methods for continuous control tasks in the following . In deep RL the parameters θ are the weights in a deep neural network function approximation of the policy or Actor as well as the state-action value function Q or Critic , respectively , and are optimized by gradient decent . The agent i.e . the Actor-Critic is trained to maximize the expected γ-discounted cumulative reward Eπ [ ∑T t=0 γ t rπ ( st , at ) ] , with respect to the policy network i.e . its parameters θ . For notational simplicity we omit the explicit dependency of the latter in the remainder of this work . Furthermore the state-action value function Q ( st , at ) , returns the value of performing a given action at while being in the state st . The Q-function is trained by minimizing the so called Bellman error as Q̂i+1 ← argmin Q E [ ( rt + γ Q̂i ( st+1 , at+1 ) −Qi ( st , at ) ) 2 ] , ( 1 ) This is commonly referred to as the ith policy evaluation step where the hat denotes the target Q-function . In offline RL one aims to learn an optimal policy for the given the dataset D =⋃ k=1 , ... , # data-points ( st , at , rt , st+1 ) k as the option for exploration of the MDP is not available . The policy is optimized to maximize the state-action value function via the policy improvement πi+1 ← argmax π Est∼D [ Q̂ ( st , πi ( at|st ) ) ] . ( 2 ) Note that behavioural policies including sub-optimal or randomized ones may be used to generate the static dataset D. In that case offline RL algorithms face difficulties in the learning process . CQL algorithm : CQL is built on top of a Soft-Actor Critic algorithm ( SAC ) ( Haarnoja et al. , 2018 ) , which employs soft-policy iteration of a stochastic policy ( Haarnoja et al. , 2017 ) . A policy entropy regularization term is added to the policy improvement step in Eq . ( 2 ) as πi+1 ← argmax π Est∼D [ Q̂ ( st , πi ( at|st ) ) − α log πi ( at|st ) ] . ( 3 ) where α either is a fixed hyperparameter or may be chosen to be trainable . CQL reduces the overestimation of state-values - in particular those out-of distribution from D. It achieves this by regularizing the Q-function in Eq . ( 1 ) by a term minimizing its values over out of distribution randomly sampled action . In the following at+1 is given by the prediction of the policy πi ( at+1|st+1 ) and α̃ is a hyperparameter balancing the regulizer term . The policy optimisation step is given by Q̂i+1 ← argmin Q ( E st , at , st+1∼D [ ( rt + γ Q̂i ( st+1 , at+1 ) −Qi ( st , at ) ) 2 ] ( 4 ) + α̃ E st∼D [ log ∑ a exp ( Qi ( st , a ) ) − E a∼π ( st ) [ Qi ( st , a ) ] ] ) 2.2 KOOPMAN THEORY . Historically , the Koopman theoretic perspective of dynamical systems was introduced to describe the evolution of measurements of Hamiltonian systems ( Koopman , 1931 ; Mezić , 2005 ) . The underlying dynamic of most modern reinforcement learning tasks is of nonlinear nature , i.e . the agents actions lead to changes of it state described by a complex non-linear dynamical system . In contrast to linear systems which are completely characterized by their spectral decomposition non-linear systems lack such a unified characterisation . The Koopman operator theoretic framework describes nonlinear dynamics via a linear infinite-dimensional Koopman operator and thus inherits certain tools applicable to linear control systems ( Mauroy et al. , 2020 ; Kaiser et al. , 2021 ) . In practice one aims to find a finite-dimensional representation of the Koopman operator which is equivalent to obtaining a coordinate transformations in which the nonlinear dynamics are approximately linear . A general non-affine control system is governed by the system of non-linear ordinary differential equations ( ODEs ) as ṡ = f ( s , a ) ( 5 ) where s , is the n-dimensional state vector , a the m-dimensional action vector with ( s , a ) ∈ S ×A the state-action-space . Moreover , ṡ = ∂ts is the time derivative , and f is some general non-linear - at least C1-differentiable - vector valued function . For a discrete time system , Eq . ( 5 ) takes the form st+1 = F ( st , at ) ( 6 ) where st denotes the state at time t where F is at least C1-differentiable vector valued function . Definition 1 ( Koopman operator ) Let K ( S ×A ) be the ( Banach ) space of all measurement functions ( observables ) . Then the Koopman operator K : K ( S ×A ) → K ( S ×A ) is defined by Kg ( st , at ) = g ( F ( st , at ) , at+1 ) = g ( st+1 , at+1 ) , ∀g ∈ K ( S ×A ) ( 7 ) where g : S ×A → R. Many systems can be modeled by a bilinearisation where the action enters the controlling equations 5 linearly as f ( s , a ) = f0 ( s ) + ∑m i fi ( s ) ai for fi , i = 0 , . . . , m i.e . C1-differentiable-vector valued functions . In that case the action of the Koopman operator takes the simple form g ( st+1 ) = K ( a ) g ( st ) = ( K0 + m∑ i Kiai ) g ( st ) , ∀g ∈ K ( S ) , ( 8 ) where K ( S ) is a ( Banach ) space of measurement functions K0 , Ki decompose the Koopman operator in a the free and forcing-term , respectively . Details on the existence of a representation as in equation 8 are discussed in Goswami & Paley ( 2017 ) . Associated with a Koopman operator is its eigenspectrum , that is , the eigenvalues λ , and the corresponding eigenfunctions φλ ( s , a ) , such that [ Kφλ ] ( s , a ) = λφλ ( s , a ) . In practice one derives a finite set of observable g⃗ = ( g1 , . . . , gN ) in which case the approximation to the Koopman operator admits a finite-dimensional matrix representation . The N ×N matrix representing the Koopman operator may be diagonalized by a matrix U containing the eigen-vectors of K as columns . In which case the eigen-functions are derived by φ⃗ = Ug⃗ and one infers from Eq . ( 5 ) that φ̇i = λiφi , for i = 1 , . . . , N with eigenvalues λi=1 , ... , N . These ODEs admit simple solutions for their time-evolution namely the exponential functions exp ( λit ) ..

Summary: The paper uses koopman theory to design principled data augmentation method for offline reinforcement learning. They learn a VAE style encoder-decoder model such that $D(E(s))=s$ and a forward model such that $F(s_t, a_t) = D((K_0 + \sum_{i=1}^m K_i, a_{t,i})E(s_t)) = s_{t+1}$. The koopman operator $K$ is then used to generate symmetries which are applied to both $s_t$ and $s_{t+1}$ as data augmentation during bellman error minimization. The resulting algorithm KFC and KFC++ leads to overall better data augmentation and improves for S4RL and CQL.