Year: 2021

Abstract:
Offline reinforcement learning leverages large datasets to train policies without interactions with the environment. The learned policies may then be deployed in real-world settings where interactions are costly or dangerous. Current algorithms over-fit to the training dataset and as a consequence perform poorly when deployed to out-of-distribution generalizations of the environment. We aim to address these limitations by learning a Koopman latent representation which allows us to infer symmetries of the system’s underlying dynamic. The latter is then utilized to extend the otherwise static offline dataset during training; this constitutes a novel data augmentation framework which reflects the system’s dynamic and is thus to be interpreted as an exploration of the environments phase space. To obtain the symmetries we employ Koopman theory in which nonlinear dynamics are represented in terms of a linear operator acting on the space of measurement functions of the system and thus symmetries of the dynamics may be inferred directly. We provide novel theoretical results on the existence and nature of symmetries relevant for control systems such as reinforcement learning settings. Moreover, we empirically evaluate our method on several benchmark offline reinforcement learning tasks and datasets including D4RL, Metaworld and Robosuite and find that by using ourframework we consistently improve the state-of-the-art for Q-learning methods.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARIES AND BACKGROUND

heading: 2.1 OFFLINE RL & CONSERVATIVE Q-LEARNING
text: 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) ]])

heading: 2.2 KOOPMAN THEORY
text: 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)..

heading: 3 THE KOOPMAN FORWARD FRAMEWORK
text: We initiate our discussion with a focus on symmetries in dynamical control systems in Subsection 3.1 where we additionally present some theoretical results. We then proceed in Subsection 3.2 by presenting the Koopman forward framework for Q-learning based on CQL. Moreover, we discuss the algorithm as well as the Koopman Forward model’s deep neural network architecture. Overview: Our theoretical results culminate in Theorem 3.4 and 3.5, which provides a road-map on how specific symmetries of the dynamical control system are to be inferred from a VAE forward prediction model. In particular, Theorem 3.4 guarantees that the procedure leads to symmetries of the system (at least locally) and Theorem 3.5 that actual new data points can be derived by applying symmetry transformations to existing ones. Practically the VAE model parametrized by a neural net is trained data from one of many behavior policies and is thus learning an approximate dynamics. The theoretical limitations are twofold, firstly the theorems only hold for dynamical systems with differentiable state-transitions; secondly, we employ a Bilinearisation Ansatz for the Koopman operator of the system. Practically many RL environments incorporate dynamics with discontinuous “contact” events where the Bilinearisation Ansatz may not be applicable. However, empirically we find that our approach nevertheless is successful for environments with “contact” events. The latter does not affect the performance significantly (see Appendix D.1).

heading: 3.1 SYMMETRIES OF DYNAMICAL CONTROL SYSTEM
text: In general, a symmetry group Σ of the state space may be any subgroup of the group of isometrics of the Euclidean space En. We restrict oneself to an Euclidean state space, however in general one may consider Riemannian manifolds, see e.g. Giannakis (2019). In particular, in our work we also consider Σ invariant compact subsets of Euclidean space. Moreover, relevant for a system of ODEs are local Lie symmetries as well as symmetries of ODEs.
Definition 2 (Local Lie Group) A parametrized set of transformations σϵ : S → S with s 7→ s̃(s, ϵ) for ϵ ∈ (ϵlow, ϵhigh) where ϵlow < 0 < ϵhigh is a one-parameter local Lie group if
1. σ0 is the identity map i.e for ϵ = 0 such that s̃(s, 0) = s.
2. σϵ1σϵ2 = σϵ2σϵ1 = σϵ1+ϵ2 for every |ϵ1|, |ϵ2| ≪ 1. 3. s̃(s, ϵ) admits a Taylor series expansion in ϵ, i.e. in a neighbourhood of s determined by
ϵ = 0 as s̃(s, ϵ) = s+ ϵ ζ(s) +O(ϵ2).
The points (1) and (2) in definition 2 imply the existence of an inverse element σϵ−1 = σ−ϵ for |ϵ| ≪ 1. Moreover, note that a system of ODEs may be represented by the set of its solutions.
Definition 3 (Symmetries of ODEs) A symmetry of a system of ODEs on a locally-smooth structure (such as En) is a locally-defined diffeomorphism that maps the set of all solutions to itself.
Definition 4 (Equivariant Dynamical System) Consider the dynamical system ṡ = f(s) and let Σ be a group acting on the state-space S. Then the system is called Σ-equivariant if f(σ · s) = σ · f(s) , for s ∈ S , ∀σ ∈ Σ. For a discrete time dynamical system st+1 = F (st) one defines equivariance analogously, namely if F (σ · st) = σ · F (st) , for st ∈ S , ∀σ ∈ Σ.
Note that as a local Lie group may satisfy the group axioms for sufficiently small parameters values it may not be a group on the entire set. And moreover not every diffeomorphism is also an isometry. Let us start by introducing symmetries in the simple context of dynamical system without control (Sinha et al., 2020) given by ṡ = f(s) where we use analog notations as in Eq. (5). The Koopman operator of equivariant dynamical system is reviewed in the appendix A. Let us next turn to the case relevant for RL, namely control systems. In the remainder of this section we focus on dynamical systems given as in Eq. (6) and Eq. (8).
Definition 5 (Action Equivariant Dynamical System) Let Σ̄ be a group acting on the state-actionspace S ×A of a general control system as in Eq. (6) such that it acts as the identity operation on A i.e. σ · (st, at) = (σ|S · st, at) , ∀σ ∈ Σ̄ . Then the system is called Σ̄-action-equivariant if
F (σ · (st, at)) = σ|S · F (st, at) , for (st, at) ∈ S ×A , ∀σ ∈ Σ̄ . (9)
Lemma 3.1 The map ∗̄ : Σ̄ × K(S × A) → K(S × A) given by (σ∗̄g)(s, a) 7−→ g(σ−1 · s, a) defines a group action on the Koopman space of observables K(S ×A).
Theorem 3.2 A Koopman operator K of a Σ̄-action-equivariant system st+1 = F (st, at) satisfies
[σ∗̄(Kg)](st, at) = [K(σ∗̄g)](st, at) . (10)
In particular, it is easy to see that the biliniarisation in Eq. (8) is Σ-action-equivariant if fi(σ|S · s) = σ|S · fi(s) , ∀i = 0, . . . ,m. Let us thus turn our focus on the relevant case of a control system st+1 = F̃ (st, at) which admits a Koopman operator description as
g(st+1) = K(at)g(st) , for at ∈ A , ∀g ∈ K(S) , (11)
where {K(a)}a∈A is a family of operators with analytical dependence on a ∈ A. Note that the bilinearisation in Eq. (8) is a special case of Eq. (11). Furthermore, let {U(a)}a∈A be a family of invertible operators s.t. U(a) : K(S) → F(S × A) is a mapping to the (Banach) space of eigenfunctions F(S×A) with eigenfunctions φ(s, a) := U(a)g(s) which obeys U(a)K(a)U(a)−1 = Λ(a), with Λ(a)φ(s, a) = λφ(a)φ(s, a) and where λφ(a) : A → R. The existence of such operators on an infinite dimensional space requires the Koopman operator in Eq. (11) to be self-adjoint or a finite-dimensional (approximate) matrix representation to be diagonalizable.1
Lemma 3.3 The map ϕ̂ : Σ̄×K(S)× {U(a)}a∈A → K(S) given by (σa∗̂g)(s) 7−→ ( U−1(a) ( σ∗̄(U(a)g) )) (s) = g(σ−1 · s) , (12)
defines a group action on the Koopman space of observables K(S). Where ∗̄ is defined analog to Lemma 3.1 but acting on F(S ×A) instead of K(S ×A) by (σ∗̄φ)(s, a) 7−→ φ(σ−1 · s, a). We refer to a control system in Eq. (11) admitting a ϕ̂-symmetry as Σ̂-action-equivariant.
Theorem 3.4 Let st+1 = F̃ (st, at) be a Σ̂-action-equivariant control system with a symmetry action as in Lemma 3.3 which furthermore admits a Koopman operator representation as
g(st+1) = K(at)g(st) , for at ∈ A , ∀g ∈ K(S) . (13)
Then [ σat ∗̂ ( K(at)g )] (st) = [ K(at) ( σat ∗̂g )] (st) . (14)
Moreover, a control system obeying equations 13 and 14 is Σ̂-action-equivariant locally if g−1 exists for a neighborhood of st, i.e. then σ · F̃ (st, at) = F̃ (σ · (st, at)).
1See Appendix A for the details. We found that, in practice the Koopman operator leaned by the neural nets is diagonalizable almost everywhere.
Let us next provide theoretical statements on how data-points may be shifted by symmetry transformations of solutions of ODEs. To establish an easier connection to the next section we introduce the notation Let E : S → K(S) and D : K(S)→ S denote the C1-differentiable encoder and decoder to and from the finite-dimensional Koopman space approximation, respectively, i.e. E◦D = D◦E = id.
Theorem 3.5 Let st+1 = F̃ (st, at) be a control system as in equation 13 and σat an operator obeying equation 14. Then σϵat : ( st, st+1, at ) 7−→ ( s̃t, s̃t+1, at ) with
s̃t = D ( ( 1 + ϵ σat ) ∗̂E(st) ) , s̃t+1 = D ( ( 1 + ϵ σat ) ∗̂E(st+1) ) (15)
is a one-parameter local Lie symmetry group of ODEs. In other words one can use a symmetry transformation to shift both st as well as st+1 such that s̃t+1 = F̃ (s̃t, at). 2

heading: 3.2 KFC ALGORITHM
text: In the previous section we laid the theoretical foundation for the KFC-algorithm by providing a raod-map on how to derive symmetries of dynamical control systems based on a Koopman latent space representation. The goal is to generate new data points for the RL algorithm at training-time as Eq. (15) in Theorem 3.5. The reward rt is not part of the symmetry shift process and will just remain unchanged as an assertion. On the power of using symmetries: Let us emphasize the practical advantage of employing symmetries to generate augmented data points. It is evident from Theorem 3.5 that a symmetry transformation shifts both st as well as st+1 which evades the necessity of forecasting states. Thus the use of an inaccurate fore-cast model is avoided and the accuracy and generalisation capabilities of the VAE are fully utilized. The magnitude of the induced shift is controlled by the parameter ϵ≪ 1 such that |s− s̃| = O(ϵ) to limit out-of-distribution generalisation errors . Algorithmically the symmetry maps are derived in two distinct ways which we denote KFC and KFC++. The latter, constitute a simple starting point to extract symmetries from our setup.3 Limitations: Theorems 3.4 and 3.5 imply that if an operator commutes with the Koopman operator we can find a local symmetry group. However, global conditions may not be easily inferred. Moreover, in practice one has that D ◦ E ≈ id. Thus all assumptions of theorem 3.5 hold approximately. The Koopman forward model: The KFC algorithm requires pre-training of a Koopman forward model F : S ×A → S which is closely related to a VAE architecture as
Fc(st, at) = D ( E(st) ) = st if c = 0: VAE D ( ( K0 + ∑m i=1Ki at,i ) E(st) ) = st+1 if c = 1: forward prediction
(16)
where both of E and D are approximated by Multi-Layer-Perceptrons (MLP’s) and the bilinear Koopman-space operator approximation are implemented by a single fully connected layer for Ki, i = 0, . . . ,m, respectively. The model is trained on batches of the offline data-set tuples (st+1, st, at) and optimized via an additive loss-function of the VAE and the forward prediction part of the model. We refer the reader to Appendix C for details.
We integrate our framework into a specific Q-learning algorithm (CQL). In practice the Koopman latent space representation is N-dimensional (finite). The Koopman operator and the symmetry generators admit matrix representations; matrix multiplication replaces the mapping ∗̂. Following in the footsteps of (Sinha et al., 2021) our approach leaves the policy improvement Eq. (3) unchanged but modifies the policy optimisation step Eq. (4) as

heading: Q̂i+1 ← argmin
text: Q
( E
st,at,st+1∼D
[( rt + γ Q̂i ( σϵat ( s̃t+1|st+1 ) , at+1 ) −Qi ( σϵat ( s̃t|st ) , at ))2] + α̃ E
st∼D
[ log ∑ a exp ( Qi(st, a) ) − E a∼π(st) [ Qi(st, a) ]]) (17)
2No assumptions on the Koopman operator are imposed. Moreover, note that the equivalent theorem holds when ϵ σat → ∑ I=1 ϵI σ I at to be a local N-parameter Lie group.
3More elaborate studies employing the extended literate on Koopman spectral analysis are desirable. See appendix B for details on numerical errors.
The state-space symmetry generating function σϵa : S → S depends on the normally distributed random variables ϵ. Note that we only modify the Bellman error of Eq. (17) and leave the CQL specific regulizer untouched. We study two distinct cases which differ on an algorithmic level4
KFC σϵat : s 7→ s̃ = D (( 1+ ϵσat ) E(s) ) KFC++ σϵat : s 7→ s̃ = D (( 1+ σat (⃗ϵ) ) E(s) ) ,
where σat (⃗ϵ) = Re ( U(at) diag(ϵ1, . . . , ϵN )U−1(at) ) , (18)
and U(at) diagonalizes the Koopman operator with eigenvalues λi , i = 1, . . . , N i.e. the latter can be expressed as K(at) = U(at)diag(λ1, . . . , λN )U−1(at). Note that we abuse our notation as E(s) ≡ g⃗(s) = [g1(s), . . . , gN (s)] i.e. the encoder provides the Koopman space observables. From theorem 3.5 one infers that in order for Eq. (18) to be a local Lie symmetry group of the approximate ODEs captured by the neural net Eq. (16) one needs σat to commute with the approximate Koopman operator in Eq. (16). For the KFC option in Eq. (18) the symmetry generator matrix σa is obtained by solving the equation σat · K(at)−K(at) · σat = 0 which may be accomplished by employing a Sylvester algorithm (syl). For KFC++ we compute the eigen-vectors of K(a) which then constitute the columns of U(a). Thus in particular one infers that [σat (⃗ϵ),K(at)] = 0. The latter, is solved by construction of σa(⃗ϵ) which commutes with the Koopman operator for all values of the random variables.5 The advantage of KFC is that it is computationally less expensive than KFC++ , however it provides less freedom to explore different symmetry directions than the latter. We employ our symmetry shift only with probability pK ; otherwise we use a random normally distributed state shift.

heading: 4 EMPIRICAL EVALUATION
text: In this section, we will first experiment with the popular D4RL benchmark commonly used for offline RL (Fu et al., 2021). The benchmark covers various different tasks such as locomotion tasks with Mujoco Gym (Brockman et al., 2016), tasks that require hierarchical planning such as antmaze, and other robotics tasks such as kitchen and adroit (Rajeswaran et al., 2017). Furthermore, similar to S4RL (Sinha et al., 2021), we perform experiments on 6 different challenging robotics tasks from MetaWorld (Yu et al., 2019) and RoboSuite (Zhu et al., 2020). We compare KFC to the baseline CQL algorithm (Kumar et al., 2020), and two best performing augmentation variants from S4RL, S4RL-N and S4RL-adv (Sinha et al., 2021). We use the exact same hyperparameters as proposed in the respective papers. Furthermore, similar to S4RL, we build KFC on top of CQL (Kumar et al., 2020) to ensure conservative Q-estimates during for policy evaluation.

heading: 4.1 D4RL BENCHMARKS
text: We present results in the benchmark D4RL test suite and report the normalized return in Table 1. We see that both KFC and KFC++ consistently outperform both the baseline CQL and S4RL across multiple tasks and data distributions. Outperforming S4RL-N and S4RL-adv on various different types of environments suggests that KFC and KFC++ fundamentally improves the data augmentation strategies discussed in S4RL. KFC-variants also improve the performance of learned agents on challenging environments such as antmaze: which requires hierarchical planning, kitchen and adroit tasks: which are sparse reward and have large action spaces. Similarly, KFC-variants also perform well on difficult data distributions such as “medium-replay”: which is a collected by simply using all the data that the policy encountered while training base SAC policy, and “-human” which is collected using human demonstrations on robotic tasks which results in a non-Markovian behaviour policy (more details can be found in the D4RL manuscript (Fu et al., 2021)). Furthermore, to our knowledge, the results for KFC++ are state-of-the-art in policy learning from D4RL datasets for most environments and data distributions.
We do note that S4RL outperforms KFC on the “-random” split of the data distributions, which is expected as KFC depends on learning a simple dynamics model of the data to use to guide the data augmentation strategy. Since the “-random” split consists of random actions in the environment, our simple model is unable to learn a useful dynamics model. For ablation studies see appendix D.
4Eq. (18) holds for both st as well as st+1 thus we have deliberately dropped the subscript. Moreover, note that in case (II) the symmetry transformations are of the form σat (⃗ϵ) ) = ∑ I=1 ϵI σ I at .
5[·, ·] denotes the commutator of two matrices. The Koopman operators eigenvalues and eignevectors generically are C-valued. However, the definition in Eq. (18) ensures that σa(⃗ϵ) is a R-valued matrix.

heading: 4.2 METAWORLD AND ROBOSUITE BENCHMARKS
text: To further test the ability of KFC, we perform additional experiments on challenging robotic tasks. Following (Sinha et al., 2021), we perform additional experiments with 4 MetaWorld environments (Yu et al., 2019) and 2 RoboSuite environments (Zhu et al., 2020). We followed the same method to
collect the data as described in Appendix F of S4RL (Sinha et al., 2021), and report the mean percent of goals reached, where the condition of reaching the goal is defined by the environment.
We report the results in Figure 2, where we see that using KFC to guide the data augmentation strategy for a base CQL agent, we are able to learn an agent that performs significantly better. Furthermore, we see that for more challenging tasks such as “push” and “door-close” in the MetaWorld, KFC++ outperforms the base CQL algorithm and the S4RL agent by a significant margin. These set of experiments further highlight the ability of KFC to guide the data augmentation strategy.

heading: 5 RELATED WORKS
text: The use of data augmentation techniques in Q-learning has been discussed recently (Laskin et al., 2020b;a; Sinha et al., 2021). In particular, our work shares strong parallels with (Sinha et al., 2021). Our modification of the policy evaluation step of the CQL algorithm (Kumar et al., 2020) is analogous to the one in Sinha et al. (2021). However, the latter randomly augments the data while our augmentation framework is based on symmetry state shifts. Regarding the connection to world models (Ha & Schmidhuber, 2018). Here a VAE is used to decode the state information while a recurrent separate neural network predicts future states. Their latent representation is not of Koopman type. Also no symmetries and data-augmentations are derived.
Algebraic symmetries of the state-action space in Markov Decision Processes (MDP) originate (Balaraman & Andrew, 2004) an were discussed recently in the context of RL in (van der Pol et al., 2020). Their goal is to preserve the essential algebraic homomorphism symmetry structure of the original MDP while finding a more compact representation. The symmetry maps considered in our work are more general and are utilized in a different way. Symmetry-based representation learning (Higgins et al., 2018) refers to the study of symmetries of the environment manifested in the latent representation. The symmetries in our case are derived form the Koopman operator not the latent representation directly. In (Caselles-Dupré et al., 2019) the authors discuss representation learning of symmetries (Higgins et al., 2018) allowing for interactions with the environment. A Forward-VAE model which is similar to our Koopman-Forward VAE model is employed. However, our approach is based on theoretical results providing a road-map to derive explicit symmetries of the dynamical systems as well as their utilisation for state-shifts.
In (Sinha et al., 2020) the authors extend the Koopman operator from a local to a global description using symmetries of the dynamics. They do not discuss action-equivariant dynamical control systems nor data augmentation. In (Salova et al., 2019) the imprint of known symmetries on the blockdiagonal Koopman space representation for non-control dynamical systems is discussed. This is close to the spirit of disentanglement (Higgins et al., 2018). Our results are on control setups and deriving symmetries. On another front, the application of Koopman theory in control or reinforcement learning has also been discussed recently. For example, Li et al. (2020) propose to use compositional Koopman operators using graph neural networks to learn dynamics that can quickly adapt to new environments of unknown physical parameters and produce control signals to achieve a specified goal. Kaiser et al. (2021) discuss the use of Koopman eigenfunction as a transformation of the state into a globally linear space where the classical control techniques is applicable. To the best of our knowledge, this paper is the first to discuss Koopman latent space for data augmentation.

heading: 6 CONCLUSIONS
text: In this work we proposed a symmetry-based data augmentation technique derived from a Koopman latent space representation. It enables a meaningful extension of offline RL datasets describing dynamical systems, i.e. further "exploration" without additional environment interactions. The approach is based on our theoretical results on symmetries of dynamical control systems and symmetry shifts of data. Both hold for systems with differentiable state transitions and with a Bilinearisation Ansatz for the Koopman operator. However, the empirical results show that the framework is successfully applicable beyond those limitations. We empirically evaluated our method on several benchmark offline reinforcement learning tasks D4RL, Metaworld and Robosuite and find that by using our framework we consistently improve the state-of-the-art of Q-learning algorithms.

heading: A KOOPMAN THEORY AND SYMMETRIES
text: This section provides some additional introductory material on Koopman theory and symmetries of dynamical systems.
Lemma A.1 The map ∗ : Σ×K(S)→ K(S) given by (σ ∗ g)(s) 7−→ g(σ−1 · s) defines a group action on the Koopman space of observables K(S).
Theorem A.2 LetK be the Koopman operator associated with a Σ-equivariant system st+1 = F (st). Then [σ ∗ (Kg)](s) = [K(σ ∗ g)](s) . (19) Theorem A.2 states that for a Σ-equivariant system any symmetry transformation commutes with the Koopman operator. For the proof see (Sinha et al., 2020).
To proceed let us re-evaluate some information from the main text of this work. Let {U(a)}a∈A be a family of invertible operators s.t. U(a) : K(S)→ F(S ×A) is a mapping to the (Banach) space of eigenfunctions φ(s, a) := U(a)g(s) ∈ F(S ×A). Which moreover obeys U(a)K(a)U(a)−1 = Λ(a), with Λ(a)φ(s, a) = λφ(a)φ(s, a) and where λφ(a) : A → R. The existence of such operators puts further restriction on the Koopman operator in Eq. (11). However, as our algorithm employs the finite-dimensional approximation of the Koopman operator i.e. its matrix representation this amounts simple for K(a) to be diagonalizable and U(a) is the matrix containing its eigen-vectors as columns. To evolve a better understanding on the required criteria for the infinite-dimensional case we employ an alternative formulation of the so called spectral theorem below.
Theorem A.3 (Spectral Theorem) Let K be a bounded self-adjoint operator on a Hilbert spaceH. Then there is a measure space (S,Σ, µ) and a real-valued essentially bounded measurable function λ on S and a unitary operator U : H → L2µ(S) i.e.U∗U = U U∗ = id such that
U ΛU∗ = K , with [ Λφ ] (s) = λ(s)φ(s) . (20)
In other words every bounded self-adjoint operator is unitarily equivalent to a multiplication operator. In contrast to the finite-dimensional case we need to slightly alter our criteria to U(a)K(a)U(a)−1 = Λ(a), with Λ(a)φ(s, a) = λφ(a, s)φ(s, a) and where λφ(s, a) : S × A → R. Concludingly, a sufficient condition for our criteria to hold in terms of operators on Hilbert spaces is that the Koopman operator K(a) is self-adjoint i.e. that
K(a) = K(a)∗ . (21)

heading: B PROOFS
text: In this section we provide the proofs of the theoretical results of Section 3.
Proof of Lemma A.1 and Theorem A.2 The proofs of the lemma as well as the theorem can be found in (Sinha et al., 2020).
Proof of Lemma 3.1: We aim to show that map ∗̄ : Σ̄ × K(S × A) → K(S × A) given by (σ∗̄g)(s, a) 7−→ g(σ−1 ·s, a) defines a group action on the Koopman space of observables K(S×A), where Σ̄ defines the symmetry group of definition 5. Firstly, let g ∈ K(S × A) and σ0 ∈ Σ̄ be the identity element, then we see that it provides existence of an identity element of ∗̄ by
(σ0∗̄g)(s, a) = g ( σ−10 · s, a ) = g(s, a) .
Secondly, let σ1, σ2 ∈ Σ̄ and ⊙ denoting the group operation i.e. σ1 ⊙ σ2 = σ3 ∈ Σ̄. Then( σ2∗̄ ( σ1∗̄g )) (s, a) = σ2∗̄ ( g ( σ−11 · s, a )) = g ( σ−12 · (σ −1 1 · s), a ) = g ( (σ−12 ⊙ σ −1 1 ) · s), a ) (I) = g ( (σ1 ⊙ σ2)−1 · s, a ) = (( σ1 ⊙ σ2 ) ∗̄g ) (s, a) = ( σ3∗̄g ) (s, a) ,
where in (I) we have used the invertibility of the group property of Σ̄. Lastly it follows analogously that for σ, σ−1 ∈ Σ̄ that(
σ∗̄ ( σ−1∗̄g )) (s, a) = (( σ−1 ⊙ σ ) ∗̄g ) (s, a) = ( σ0∗̄g ) (s, a) = g(s, a) .
Thus the existence of an inverse is established which concludes to show the group property of ∗̄.
Proof of Theorem 3.2: We aim to show that with K being the Koopman operator associated with a Σ̄-action-equivariant system st+1 = F (st, at). Then[
σ∗̄ ( Kg ) ] (s, a) = [ K ( σ∗̄g )] (s, a) .
First of all note that by the definition of the Koopman operator of non-affine control systems it obeys
[Kg](st, at) = g(F (st, at), at+1) .
Using the latter one thus infers that[ σ∗̄ ( Kg )] (st, at) = σ∗̄g(F (st, at), at+1) = g(σ−1 · F (st, at), at+1) (I) = g(F (σ−1 · st, at), at+1) ,
where in (I) we have used that it is a Σ̄-action-equivariant system. Moreover, one finds that g(F (σ−1 · st, at), at+1) = Kg(σ−1 · st, at) = [ K ( σ∗̄g ) ] (st, at) ,
which concludes the proof.
Proof of Lemma 3.3: We aim to show that the map ϕ̂ : Σ̄×K(S)× {U(a)}a∈A → K(S) given by (
σa∗̂g ) (s) 7−→ ( U−1(a) ( σ∗̄(U(a)g) )) (s) = g ( σ−1 · s ) ,
defines a group action on the Koopman space of observables K(S). Where ∗̄ is defined analog to Lemma 3.1 but acting on F(S ×A) instead of K(S ×A) by (σ∗̄φ)(s, a) 7−→ φ(σ−1 · s, a). First of all note that(
U−1(a) ( σ̄ ∗ (U(a)g) )) (s) = U−1(a) ( σ̄ ∗ φ(s, a) ) = U−1(a)φ ( σ−1 · s, a ) = g ( σ−1 · s ) (22)
We proceed analogously as in the proof of Lemma 3.1 above. Firstly, let g ∈ K(S) and σ0 ∈ Σ̄ be the identity element then on infers from Eq. (22) that it provides existence of an identity element of ϕ̂ by
(σa,0∗̂g)(s) = g ( σ−10 · s ) = g(s) ,
where we have used the notation (σa,i∗̂g)(s) = ( U−1(a) ( σi∗̄(U(a)g) )) (s) , for i = 0, . . . .
Secondly, let σ1, σ2 ∈ Σ̄ and ⊙ denoting the group operation i.e. σ1 ⊙ σ2 = σ3 ∈ Σ̄. Then( σa,2∗̂ ( σa,1∗̂g )) (s) = σa,2∗̂ ( g ( σ−11 · s )) = g ( σ−12 · (σ −1 1 · s) ) = g ( (σ−12 ⊙ σ −1 1 ) · s ) (I) = g ( (σ1 ⊙ σ2)−1 · s ) = g ( σ−13 · s ) = ( σa,3∗̄g ) (s) ,
where in (I) we have used the invertibility of the group property of Σ. Lastly, it follows analogously that for σ, σ−1 ∈ Σ̄ that(
σ∗̂ ( σ−1∗̂g )) (s) = (( σ−1 ⊙ σ ) ∗̂g ) (s) = ( σ0∗̂g ) (s) = g(s) .
Thus the existence of an inverse is established which concludes to show the group property of ϕ̂.
Proof of Theorem 3.4: We aim to show twofold.
=⇒: Firstly, that with st+1 = F̃ (st, at) be a Σ̂-action-equivariant control system with a symmetry action as in Lemma 3.3 which furthermore admits a Koopman operator representation as
g(st+1) = K(at)g(st) , for at ∈ A , ∀g ∈ K(S) .
Then [ σat ∗̂ ( K(at)g )] (st) = [ K(at) ( σat ∗̂g )] (st) .
⇐=: Secondly, the converse. Namely, if a control system st+1 = F̃ (st, at) obeys Eqs. (13) and (14), then it is Σ̂-action-equivariant, i.e. σ · F̃ (st, at) = F̃ (σ · (st, at)). For notational simplicity we drop the subscripts referring to time in the remainder of this proof i.e. st → s and at → a. Let us start with the first implication i.e. =⇒. First of all note that by the definition of the Koopman operator one has[
K(a)g ] (s) = g ( F̃ (s, a) ) .
Using the latter one infers that[ σa∗̂ ( K(a)g )] (s) = σa∗̂g(F̃ (s, a)) ,
= U−1(a) ( σ∗̄ ( U(a)g ( F̃ (s, a) ))) , = U−1(a) ( σ∗̄φ ( F̃ (s, a), a )) , = U−1(a) ( φ ( σ−1 · F̃ (s, a), a ))) , = U−1(a) ( φ ( F̃ ( σ−1 · s, a ) , a ))) , = g ( F̃ ( σ−1s, a )) .
Moreover, one derives that g ( F̃ ( σ−1s, a )) = K(a)g ( σ−1s ) ,
= K(a) U−1(a)U(a)︸ ︷︷ ︸ =id
g ( σ−1s ) ,
= K(a)U−1(a)φ ( σ−1s, a ) ,
= K(a) ( U−1(a) ( σ∗̄φ ( s, a ))) , = K(a) ( U−1(a) ( σ∗̄ ( U(a)g(s) ))) , = [ K(a) ( σa∗̂g )] (s) ,
which concludes the proof of the first part of the theorem. Let us next show the converse implication i.e. ⇐=. For this case it is practical to use the discrete time-system notation explicitly. Let σ ∈ Σ be a symmetry of the state-space and be s̃t = σ · st and s̃t+1 = σ · st+1 the σ-shifted states. Then
g ( st+1 ) = g ( F̃ (st, at) ) ,
= K(at) g(st) , = K(at) g ( σ−1 · s̃t ) ,
= K(at) ( σa∗̂g ( s̃t )) , = [ K(at) ( σa∗̂g )] (s̃t) ,
(I) = [ σa∗̂ ( K(at)g )] (s̃t) .
Thus in particular σ−1a ∗̂g ( st+1 ) = [ σ−1a ∗̂ ( σa∗̂ ( K(at)g ))] (s̃t) Lemma 3.3 = [ K(at)g ] (s̃t) ,
where in (I) we have used that the symmetry operator commutes with the Koopman operator. Moreover, one finds that
σ−1a ∗̂g ( st+1 ) = g (( σ−1 )−1 · st+1) = g(σ · st+1) = g(s̃t+1) , from which one concludes that
g(s̃t+1) = [ K(at)g ] (s̃t) = g ( F̃ (s̃t, at) ) .
Finally, we use the invertibility of the Koopman space observables i.e. g−1(g(s)) = s to infer
s̃t+1 = F̃ (s̃t, at) .
However, in general s the Koopman space observables are not invertible globally. The "inverse function theorem" guarantees the existence of a local inverse if g(s) is C1 differentiable for maps between manifolds of equal dimensions. However, we assume inevitability locally.
s̃t+1 = σ · st+1 ⇒ F̃ (s̃t, at) = σ · F̃ (st, at) ⇒ F̃ (σ · st, at) = σ · F̃ (st, at) , (23)
which at last concludes our proof of the second part of the theorem.
Extension of Theorem 3.5: Moreover, one may account for practical limitations i.e. an error by the assumption [σϵat ,K(at)] = ϵa 1. One then finds that s̃t+1 = F̃ (s̃t, at) +O(ϵa). Thus the error becomes suppressed when ϵa ≪ ϵ. The error may be due to practical limitations of capturing the true dynamics as well as the symmetry map.
Proof of Theorem 3.5: We will show here theorem 3.5 as well as the extension mentioned in the paragraph prior to this proof. Let st+1 = F̃ (st, at) be a Σ̂-action-equivariant control system as in Theorem 3.4. For symmetry maps
s̃t = D ( ( 1 + ϵ σat ) ∗̂E(st) ) , s̃t+1 = D ( ( 1 + ϵ σat ) ∗̂E(st+1) ) . (24)
we aim to show that one can use a symmetry transformation to shift both st as well as st+1 σϵat : ( st, st+1, at ) 7−→ ( s̃t, s̃t+1, at ) , s.t. s̃t+1 = F̃ (s̃t, at) . (25)
From equation 24 and the definition of the Koopman operator one infers that6 s̃t+1 = D ( σat ∗̂E(st+1) ) = D ( σat ∗̂K(at)E(st) ) . (26)
By using that [σat ,K(at)] = 0 and equation 24 one finds that s̃t+1 = D ( K(at)σat ∗̂E(st) ) = D ( K(at)E(D(︸ ︷︷ ︸
=1
σat ∗̂E(st))) ) = D ( K(at)E(s̃t) ) , (27)
which concludes the proof of the first part of the theorem with σat → (1 + ϵ σat). Local diffeomorphism. Per definition the maps D,E are differentiable and invertible which implies that they provide a local diffeomorphism from and to Koopman-space. Also the linear map i.e. a matrix multiplication by (1 + ϵ σat) is a diffeomorphism. Thus the symmetry map σ ϵ at i.e. Eq. (15) constitutes a local diffeomorphism. The above proof culminating in Eq. (24) implies that we have a local diffeomorphism mapping solutions of the ODEs to solutions, thus a symmetry of the system of ODEs. Limitations. However D,E to be invertible is a strong assumption as it implies that the Koopman space approximation admit the same dimension as the state space. Note that however in practice as we only require D ◦ E ≈ id one may choose other Koopman space dimensions. Local Lie group. What is left to show is that the symmetry of the system of ODEs locally is a Lie group. In definition 2 we need to show points (1)-(3).
1. For ϵ = 0 one finds that σ0at is the identity map i.e for such that s̃t(st, 0) = D ( (1 + 0 · σat)E(st) ) = D(E(st)) = st (28)
6For notational simplicity we study the general case (1 + ϵ σat) → σat .
2.
σϵ1σϵ2 = D ( (1+ ϵ1σat)(1+ ϵ2σat)E(st) ) (29)
= D ( (1+ ϵ1σat + ϵ2σat +O(ϵ1 · ϵ2))E(st) ) = D ( (1+ (ϵ1 + ϵ2)σat +O(ϵ1 · ϵ2))E(st)
) = σϵ1+ϵ2
for every |ϵ1|, |ϵ2| ≪ 1. 3. s̃t(st, ϵ) admits a Taylor series expansion in ϵ, i.e. in a neighbourhood of s determined by
ϵ = 0. We may Taylor expand D around the point E(st) ≡ g as
s̃t(st, ϵ) = D(E(st)) + ϵ N∑ I=1 (σatE(st))I ∂D ∂gI +O(ϵ2) (30)
thus s̃t(st, ϵ) = st + ϵζ(st) +O(ϵ2) (31)
for ζ(st) = ∑N
I=1(σatE(st))I ∂D ∂gI .
4. The existence of an inverse element (σϵat) −1 = σ−ϵat follows from (1) and (2).
Numerical errors. Let us next turn to the second part of the theorem incorporating for practical numerical errors. We aim to show that under the assumption [σϵat ,K(at)] = ϵa 1 one finds s̃t+1 = F̃ (s̃t, at) +O(ϵa). Note that σat → (1 + ϵ σat) admits an expansion in the parameter ϵ. For ϵ≪ 1 one can Taylor expand the differentiable functions D,E to find that
s̃t = st + δ(ϵ) , with δ(ϵ) = ϵ ∑ I gσI ∂D ∂gI ∣∣ E(st) +O(ϵ2) , (32)
where gσ = σat ∗̂E(st) and g = E(st), and with the index I = 1, . . . , N . The analog expression holds for t+ 1 i.e. s̃t+1 = st+1 +O(ϵ). If at = at(st) is differentiable function of the states7 one may Taylor expand the Koopman operator using equation 32 as8
K(at(s̃t)) = K(at(st)) + ∑ i ∑ n δn(ϵ) a ′ in(st)Ki = K(at(st)) + ϵ∆+O(ϵ2) . (33)
where and a′in(st) = ∂at ∂st,n |st and n = 1, . . . , dim(st) and with
∆ = ∑ i ∑ n ∑ I gσI ∂Dn ∂gI ∣∣ E(st) a′in(st)Ki . (34)
However, by the definition of the symmetry map at(s̃t) = at(st), thus the Koopman operators in equation 33 match ∀ϵ. By using the assumption on the error one infers from equation 26 that
s̃t+1 = D ( (1+ ϵσat)∗̂K(at)E(st) ) = D (( K(at) + ϵa 1 ) (1+ ϵσat)∗̂E(st) ) = D ( K(at)) (1+ ϵσat)∗̂E(st) + ϵa E(st) +O(ϵaϵ)
) = D ( K(at)E(s̃t) + ϵa E(st) +O(ϵaϵ) ) , (35)
One may Taylor expand D by making use of ϵa ≪ 1 as
s̃t+1 = D ( K(at)E(s̃t) ) + ϵa ∑ I EI(st) ∂D ∂gI ∣∣∣ K(at)E(s̃t) +O(ϵaϵ) +O(ϵ2a) . (36)
7Although neural networks modeling the policies modeling the data-distribution contain non-differentiable components they are differentiable almost everywhere.
8The reader may wonder about the explicit use of the index δn(ϵ) expressing the dimension of the state. it is necessary here although we have simply used st as an implicit vector without explicitly specifying any components throughout the work.
Thus under the assumption that the numerical error is given by a violation of the commutation relation one finds s̃t+1 = F̃ (s̃t, at) +O(ϵa) where
ϵa ∑ I EI(st) ∂D ∂gI ∣∣∣ K(at)E(s̃t) ∼ O(ϵa) . (37)
By comparison with equation 32 one infers that the ϵ-expansion is a good approximation to the real dynamic if ϵa ≪ ϵ. This concludes the proof of theorem 3.5. Let us note that the proof for state shifts arising from a sum of symmetry transformations as
s̃t = D ( ( 1 + N∑ I=1 ϵI σ I at ) ∗̂E(st) ) , s̃t+1 = D ( ( 1 + N∑ I=1 ϵI σ I at ) ∗̂E(st+1) ) . (38)
is completely analogous, and an analogous theorem holds. We choose the number of symmetries to run over the dimension of the Koopman latent space N as it is relevant for the KFC++ setup.
C IMPLEMENTATION DETAILS
The Koopman forward model The KFC algorithm requires pre-training of a Koopman forward model F : S ×A → S as
Fc(st, at) = D ( E(st) ) = st if c = 0 corresponds to a VAE. D ( ( K0 + ∑m i=1Ki at,i ) E(st) ) = st+1 if c = 1: forward prediction model.
(39) where both of E and D are approximated by Multi-Layer-Perceptrons (MLP’s) and the bilinear Koopman-space operator approximation are implemented by a single fully connected layer for Ki=0,...,m, respectively. The model is trained on batches of the offline data-set tuples (st+1, st, at) and optimized via the loss-function
l̄2(F1(st, at) , st+1) + γ2 l̄2(F0(st + sϵ, at) , st + sϵ) , (40)
where sϵ is a normally distributed random sate vector shift with zero mean and a standard deviation of 6 · 10−2 where the l̄2-loss is the Huberloss. For training the Koopman forward model in Eq. (16) we split the dataset in a randomly selected training and validation set with ratios 70%/30%. We then loop over the dataset to generate the required symmetry transformations and extend the replay buffer dataset with the latter according to Eq. (18) as
KFC : ( st, st+1, at, rt ) symmetries7−→ (st, st+1, at, rt, σat) , (41) KFC++ : ( st, st+1, at, rt
) symmetries7−→ (st, st+1, at, rt, U(at), U(at)−1) . It is more economic to compute the symmetry transformation as a pre-processing step rather than at runtime of the RL algorithm.
Let us reemphasize a couple of points addressed already in the main text to make this section selfcontained. Note that we have built upon the implementation of CQL (Kumar et al., 2020), which is based on SAC (Haarnoja et al., 2018). We also use the same hyperparameters for all the experiments as presented in the CQL and S4RL papers for baselines.
Koopman forward VAE-model: For the encoder as well as the decoder we choose a three layer MLP architecture, with hidden-dimensions [512, 512, N ] and [512, 512,m], respectively. Where N = 32 is the Koopman latent space dimension and m is the state space dimensions following our notation from section 2. The activation functions are ReLU (Nair & Hinton, 2010). We use an ADAM optimizer with learning rate 3 · 10−4. We train the model for 75 epochs. The batch size is chosen to be 256. Lastly, we employ a random normally distributed state-augmentation for the VAE training i.e. s→ s+ sϵ with sϵ ∈ N (0, 1 · 10−2). Hyper-parameters KFC: We take over the hyper-parameter settings from the CQL paper (Kumar et al., 2020) except for the fact that we do not use automatic entropy tuning of the policy optimisation step Eq. (3) but instead a fixed value of α = 0.2. The remaining hyper-parameters of the conservative
Q-learning algorithm are as follows: γ = 0.99 and τ = 5 · 10−3 for the discount factor and target smoothing coefficient, respectively. Moreover, the policy learning rate is 1 · 10−4 and the value function learning rate is 3 · 10−4 for the ADAM optimizers. We enable Lagrange training for α̃ with threshold of 10.0 and learning rate of 3 · 10−4 where the optimisation is done by an ADAM optimizer. The minimum Q-weight is set to 10.0 and the number of random samples of the Q-minimizer is 10. For more specifics on these quantities see (Kumar et al., 2020). The model architectures of the policy as well as the value function are three layer MLP’s with hidden dimension 256 and ReLU -activation. The batch size is chosen to be 256. Moreover, the algorithm performs behavioral cloning for the first 40k training-steps, i.e. time-steps in an online RL notation.
The KFC specific choices are as follows: the split of random to Koopman-symmetry state shifts is 20/80 i.e. pK = 80%, whereas the option for the symmetry-type is either "Eigenspace" or "Sylvester" referring to case (I) and (II) of Eq. (18), respectively. Regarding the random variables in the symmetry transformation generation process, we choose ϵi ∈ N (0, 1 · 10−4) for case (II) and ϵ ∈ N (0, 5 · 10−5) for case (I) in Eq. (18) after normalizing σa by its mean. While for the random shift we use ϵ̃ ∈ N (0, 3 · 10−3) which is the hyper-parameter used in (Sinha et al., 2021). Computational setup: We performed the empirical experiments on a system with PyTorch 1.9.0a (Paszke et al., 2019) The hardware was as follows: NVIDIA DGX-2 with 16 V100 GPUs and 96 cores of Intel(R) Xeon(R) Platinum 8168 CPUs and NVIDIA DGX-1 with 8 A100 GPUs with 80 cores of Intel(R) Xeon(R) E5-2698 v4 CPUs. The models are trained on a single V100 or A100 GPU.

heading: D ABLATION STUDY

heading: D.1 TRAINING ON NON CONTACT EVENTS
text: To test the theoretical limitations in a practical setting one interesting experiment is to use the S4RLN augmentation scheme when the current state is a “contact” state and to choose KFC++ otherwise. We define a “contact” state as one where the agent makes contact with another surface; in practice we looked at the state dimensions which characterize the velocity of the agent parts, and if there is a sign difference between st and st+1, then contact with a surface must have been made. The details regarding which dimension maps to agent limb velocities is available in the official implementation of the Open AI Gym suite (Brockman et al., 2016).
Note that the theoretical limitations are twofold, firstly the theorems 3.4 and 3.5 only hold for dynamical systems with differentiable state-transitions; secondly, we employ a Bilinearisation Ansatz for the Koopman operator. Contact events lead to non-continuous state transitions and moreover
would require incorporating for external forces in our Koopman description, thus both theoretical assumptions are practically violated in our main empirical evaluations of KFC and KFC++. By performing this ablation study which is more suitable for our Ansatz, i.e. the KFC++ part is only trained on "differentiable" trajectories without contact events. In other words, the Koopman description is unlikely to give good information regarding st+1 if there is a “contact event”, we simply use S4RL-N when training the policies, which we expect to give slightly better performance than the KFC++ baseline.
We report the results in Table 2, where we perform results on the OpenAI Gym subset of the D4RL benchmark test suite. We denote the proposed scheme as KFC++-contact and compare to both S4RL-N and KFC++. We see that KFC++-contact performs comparably and slightly better in some instances, compared to KFC++. This is expected as contact events may be rare, or the data augmentation given by the KFC++ model during contact events is pseudo-random and itself is normally distributed, which may suggest that KFC++ behaves similarly to KFC++-contact. However, we do not say this conclusively.

heading: D.2 MAGNITUDE OF KFC AUGMENTATION
text: To further try to understand the effect of using KFC for guiding data augmentation, we measure the mean L2-norm of the strength of the augmentation during training. We measured the mean L2-norm between the original state and the augmented state, and found the distance to be 2.6 × 10−4. For comparison, S4RL uses a fixed magnitude of 1× 10−4 (Sinha et al., 2021) for all experiments with S4RL-Adv. Since S4RL-Adv does a gradient ascent step towards the direction with the highest change, its reasonable for the magnitude of the step to be smaller in comparison. An adversarial step too large may be detrimental to training as the step can be too strong. KFC minimizes the need for making such a hyperparameter search decision.

heading: D.3 SYMMETRIES
text: In this section we provide a quantitative analysis of the symmetry shifts generated by KFC++ in comparison to random shifts of the latent states. For the latter we choose the variance such that absolute state shift ∆S := |s̃t − st|+ |s̃t+1 − st+1| (42) averaged over all sampled data points (D4RL) is comparable between the two setups. Moreover we simply use the encoder We then compare how well the shifted states are in alignment with the actual dynamic using the online Mujoco environment .9 To do so we set the environment state to the value s̃t and use the action at to perform the environment step which we denote by st+1 = M(st, at)10 The performance metric is then given by
∆E := |st+1 −M(s̃t, at)| . (43) Note that ∆E is simply the error to the actual dynamic of the environment. We take the model trained on the hopper-medium-expert dataset from section appendix D.1 to compute the symmetry shift. Moreover, for the comparison in equation 43 we split the state space into positions and velocities of the state vector, i.e. the first five elements of st+1 −M(s̃t, at) give position while the remaining ones given velocities. The norm is taken subsequently. For the results see figure D.3.
We conclude that there is a qualitative difference between the distributions obtained. Most notable the symmetry shift allows for a controlled shift by a larger ∆S while still maintaining a small error ∆E. This amounts to a wide but accurate exploration of the environments state-space. Although we cannot say this conclusively we expect this to be the reason for the performance gains of KFC++ over random shifts.

heading: D.4 PREDICTION MODEL
text: As a further ablation study it is of interest to compare the symmetry induced state shifts tuple (s̃t, s̃t+1) to one obtained by the forward prediction of our VAE-forward model equation 16.
9Note that this requires a minor modification of the environment code by a function which allows to set states to specific values.
10M(st, at) denotes the step done by the active gym environment.
KFC++prediction: In particular we use the forward prediction on s̃t obtained by the KFC++ shift, i.e. s̃t+1 = Fc=1(s̃t, at). We refer to this setup in the following as KFC++prediction. See table 3 for the results. Fwd-prediction: Moreover, we study the case where we s̃t is obtained by a shift with a normally distributed random variable N (0, 6 · 10−3) and then simply forward predict as s̃t+1 = Fc=1(s̃t, at). This is comparable to conventional model based approaches employing our simple VAE forward model. This constitutes an ideal systematic comparison as the model, hyper-parameters and training procedure are identical to the one used in the KFC variants. Moreover, the variance of the random variable is such that the distance to the augmented states is comparable to the ones in KFC, see appendix D.2. See table 4 for the results.
We conclude that KFC++prediction falls behind both KFC as well as KFC++. This was expected as the symmetry shift in the latter alters the original data points by means of the VAE which admits not only a much higher accuracy but also advanced generalisation capabilities to out of distribution values. More interesting however is that Fwd-prediction falls behind both KFC as well as KFC++ as well as KFC++prediction. This is a strong indicator that the symmetry shifts provide superior out-of-distribution data for training a Q-learning algorithm.

heading: D.5 DISCUSSION
text: Although our KFC framework is suited best for the description of environments described by control dynamical systems the learned (in self-supervised manner) linear Koopman latent space representation may be applicable to a much wider set of tasks. However, there are notable shortcomings to the current implementation both conceptually as well as practically. The bilinearisation in Eq. (16) of the latent space theoretically assumes the dynamical system to be governed by Eq. (8), which is rather restrictive. Although a general non-affine control system admits a bilinearisation (Brunton et al., 2021)

heading: Task Name KFC KFC++ KFC++-prediction

heading: Domain Task Name KFC KFC++ Fwd-prediction
text: it generically requires the observables to depend on the action-space variables implicitly. Secondly, the Koopman operator formalism is theoretically defined by its action on an infinite dimensional
observable space. The finite-dimensional approximation i.e. the latent space representation of the Koopman forward model in Eq. (16) lacks accuracy due to that. On the practical side our formalism requires data pre-processing which is computationally expensive, i.e. solving the Sylvester or eigenvalue problem for every data-point. Moreover, the Koopman forward model in Eq. (16) serves as a function approximation to two distinct tasks. Thus one faces a twofold accuracy vs. over-estimation problem which needs to be balanced. The systematic error in the VAE directly imprints itself on the state-data shift in Eqs. (15) and (18) and may thus conceal any potential benefit of the symmetry considerations. Lastly, the dynamical symmetries do not infer the reward. Thus the underlying working assumption is that the reward should not vary much in Eq. (15).
A note on the simplicity of the current algorithm: Let us stress a crucial point. Algorithmically the symmetry maps are derived in two distinct ways KFC and KFC++. The latter, constitute a simple starting point to extract symmetries from our setup. More elaborate studies employing the extended literate on Koopman spectral analysis are desirable. Moreover, it is desirable to extend our framework to more complex latent space descriptions such as e.g. world models. Both on a theoretical as well as a practical level. It is our opinion that by doing so there is significant room for improvement both in the accuracy of the derived symmetry transformations as well their induced performance gains of Q-learning algorithms. Note that currently our VAE model is of very simple nature and the symmetries are extracted in a rather uneducated way. While the Sylvester algorithm simply converges to one out of many symmetry transformations for the KFC++ algorithm we omit all the information of the imaginary part, let alone utilize concrete spectral information.

Review summaries:
This paper develops a new method (Koopman Forward (Conservative) Q-learning, KFC in short) for offline reinforcement learning by extending the static dataset to include new states that are computed by learning an action-invariant Koopman latent representation of the system. The Koopman transformation to obtain the transformed state transitions is build using a variational auto-encoder. Experimental results are provided on D4RL, RoboSuite and MetaWorld benchmarks. This is a well-written paper (there are some minor typos which the authors should rectify). The method developed here shows improved performance over current methods in a number of benchmark problems.

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.