Koopman Q-learning: Offline Reinforcement Learning via Symmetries of Dynamics
Authors: Matthias Weissenbacher
Samarth Sinha
Animesh Garg
Yoshinobu Kawahara
Venue: ICLR 2022 Submitted (2022)
Keywords: offline reinforcement learning
representation learning
Koopman theory
symmetries of representations
data augmentation

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 our framework we consistently improve the state-of-the-art for Q-learning methods.

Under review as a conference paper at ICLR 2022
KOOPMAN Q-LEARNING: OFFLINE REINFORCEMENT LEARNING VIA AAASYYA SYMMETRIES OF DYNAMICS
Anonymous authors Paper under double-blind review
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 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 '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 's dynamic and is thus to be interpreted as an exploration of the environments phase space. To obtain the sym- metries we employ Koopman theory in which nonlinear dynamics 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.
INTRODUCTION
The recent impressive advances in reinforcement learning (RI.) range from robotics, to strategy games and recommendation systems (Kalashnikov et al., 2018; 18 Li al., 2010). Reinforcement learning canonically regarded as active learning process also referred to as online RI. where the agent interacts with the environment at each training run. 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 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 be collected once which alleviates the before-mentionel 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 finetuned the desired environment. Offline RL has successfully been applied to learn agents that outperform the behavioral policy used to collect the data (Kumar al., 2020; Wu et al., 2019; Agarwal al., 2020b; Ernst et al., 2005). However algorithms admit major shortcomings in regard to over werr-fiting and overestimating the true state-action values of the distribution. One solution was recently propsed by Sinha 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 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. was argued previously that is basically impossible for an offline RL agent learn optimal policy the generalization near data generically leads compounding

---

Under review as a conference paper at ICLR 2022
errors such as overestimation bias (Kumar et al., 2020). In this paper, we look offline RL through the lens of Koopman spectral theory in which nonlinear dynamics are represented in terms of a linear operator acting 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 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 's underlying dynamic in self--uuprviiee fashion from the static dataset, and second employing the latter to extend the offline dataset at training time by outto--diitributionn values. As this reflects the system's 's dynamics the additional data is to be interpreted as an exploration of the environment's '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 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, 8141) + (5c, $141) 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 neighbourhood around a MDP tuple in the offline dataset in which the behavioral policy is likely to choose the same action based 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-Jearning 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), Meta World (Yu et al., 2019) and Robosuite (Zhu et al., 2020) and find that by using KFC we can improve the state-of-the-ar 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,p. p. r, 7), with state space S, action space A, and p(s a.) the transition density function from the current state and action to the next state. Moreover, 7 is the discount factor and r(S,) the reward function. At any discrete time the agent chooses an action az E A according to its underlying policy Te (a) based on the information of the current state E S where the policy is parametrized by 0. We focus on the Actor-Critic methods for continuous control tasks in the following. In deep RL the parameters 0 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 y-discounted cumulative reward E, 1o ( at)], with respect to the policy network i.e. its parameters A. For notational simplicity we omit the explicit dependency of the latter in the remainder of this work. Furthermore the aattonction value function Q(se, at), returns the value of performing given action a, while being in the state st. The Q-function is trained by minimizing the so called Bellman error as
Qiti f- arg min E (r+ + (5(411))444))
()()))))))
Q
This is commonly referred to as the jth 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 U Bbemt,....ddatt-ppiin (St, ae Tt, st-1); 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
T++ < arg max Eg ~D Q( T (a(ls))))]
(2)
2

---

Under review as a conference paper at ICLR 2022
Learn Koopman Latent Space Representation
Apply Symmetry Shifts to Offline RL Data
8+++1
1
St
St
c3
((a+)
8(+1
1
Kooprin operator
St
Derive Symmetries of Dynamics
Input in Q-Learning Algorithm
randomly symmetry shifted states in é- ball
o symmetry
commutes with K(ttt)
policy oplimisation step
Figure 1: Overview of Koopman Q-learning. The states of the data point 8141, as) are shifted along symmetry trajectories parametrized by f1.2.3 for constant a, Symmetry transformations can be combined to reach other specific subset of the e-ball region.
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 Soft-Actor Critic algorithm (SAC) (Haarnoja et al., 2018), which employs soft-policy iteration of stochastic policy (Haarnoja et al., 2017). A policy entropy regularization term is added to the policy improvement step in Eq. (2) as
T1+11 arg max E..-D ((t,, (a(()s))) - a log (=(ees)))
where either is fixed hyperparameter 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 a++1 is given by the prediction of the policy ग. (0:+1 (s++1) and ā is hyperparameter balancing the regulizer term. The policy optimisation step is given by
rr Qi(s))1(11 ac+1) - Qi an)) )
Q411 f arg min Q
E Re
(4)
+
E Re
log Σ exp (Qi(s),
E ama(s))
[Q((6)))1]]
2.2 KOOPMAN THEORY
Historically, the Koopman theoretic perspective of dynamical systems introduced describe the evolution of measurements of Hamiltonian systems (Koopman, 1931; Mezić, 2005). The underlying dynamic of most modern reinforcement learning tasks of nonlinear nature, i.c. the agents actions lead to changes of it state described by complex non- 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 Koopman operator and thus inherits certain tools applicable
dynamics via linear 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 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)
(5) where is the n- n-dimensional state vector, a the m- nn-dimensional action vector with (s, a) E S x A the actte-action-spcee. Moreover, 5 0,s is the time derivative, and is some general non-linear at least C' differeetiiabe vector valued function. For a discrete time system, Eq. (5) takes the form F(st, 0z)
8= f(, a)
8t41
where s; denotes the state at time i where F is at least C'. -iiiifferentiable vector valued function.
Definition (Koopman operator) Let K(S A) be the (Banach) space of all measurement func- tions (observables). Then the Koopman operator K. K(S x A) + K(S x A) is defined by Kg(St, 0t) = g( at), a141) = g (8441,, G(41), Vi E K(S x A)
(7)
where g : A TR.
3

---

Under review as a conference paper at ICLR 2022
Many systems be modeled by bilinearisation where the action enters the controlling equations linearly as f(s, a) = fu(s) + E," f((s)a; a); for fs. i = O...., m i.c. C' -differentiable-veeee valued functions. In that case the action of the Koopman operator takes the simple form
g(s(+1) K(())((s)) = (Ko + IKaa))o(s)) Vg X(S)
(8)
where (S) is (Banach) space of measurement functions Kn, KG decompose the Koopman operator in a the free and forcing-term, respectively. Details the existence of representation in equation 8 are discussed in Goswami & Paley (2017). Associated with a Koopman operator is its and the corresponding eigenfunctions (s, a), such that
eigenspectrum, that is, the eigenvalues [K a) = X (s, a). In practice one derives a finite set of observable ğ = (91,... INN in which case the approximation The N x N matrix representing the Koopman operator may be diagonalized by matrix U containing the eigen-vectors of as columns. which case eigen-functions are derived by 전 Uğ and one infers from Eq. (5) that cp: Appi for i = l..... N with eigenvalues Ai= Ii......, N. These ODEs admit simple solutions for their time-evolution namely the exponential functions exp(A,t)..
the Koopman operator admits
matrix representation.
3 THE KOOPMAN FORWARD FRAMEWORK
We initiate our discussion with focus 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-Jearning 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 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 is trained data from of many behavior policies and is thus learning approximate dynamics. The theoretical limitations are twofold, firstly the theorems only hold for dynamical systems with differentiable state- taae-trranstions; secondly, employ 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 find that our approach nevertheless successful for environments with "contact" events. The latter does not affect the performance significantly (see Appendix D.
3.1 SYMMETRIES OF DYNAMICAL CONTROL SYSTEM
In general, a symmetry group Σ of the state space may be any subgroup of the group of isometrics of the Euclidean space E". We restrict oneself to 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 system of ODEs are local Lie symmetries as well as symmetries of ODEs.
Definition 2 (Local Lie Group) A parametrized set of transformations of S > S with s +> 8(s, c)
for c (Clows higgh) where Clow 0 Chigh is one-parameter local Lie group if
1. d o the identity map i.e for e = such that 3(s,0)
2. o = 0=2001 te for every |611, |c2l < 1.
3. 3(s,e) admits Taylor series expansion e, i.e. in a neighbourhood of determined by = 0 as 5(8, +) = + ((s) + O(e²).
The points (1) and (2) in definition 2 imply the existence of inverse element f-1 o for le| < 1. Moreover, note that a system of ODEs may be represented by the set of its solutions.
Definition (Symmetries of ODEs) A symmetry of a system of ODEs on a Locally-smood structure (such E.)) ) locally-defindd diffeomorphism that maps the all solutions itself.
4

---

Under review as a conference paper at ICLR 2022
Definition 4 (Equivariant Dynamical System) Consider the dynamical system à f(s) and let be a group acting on the state-space Then the system is called E-equivariant if f(o .s)
a. f(s), for E S Vo Σ. For a discrete time dynamical system 8++1 = F(s,) one defines equivariance analogously, namely if Fi 8t) =o F(+)). for St E vo E Σ.
Note that as a local Lie group may satisfy the group axioms for sufficiently small parameters values it may not be 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 á = f(s) where we use analog notations as Eq. (5). The Koopman operator of equivariant dynamical system is reviewed in the appendix A. Let next turn to the case relevant for RL, namely control systems. In the remainder of this section we focus on dynamical systems given in Eq. (6) and Eq. (8).
Definition 5 (Action Equivariant Dynamical System) Let E be group acting on the state-action- space x A of a general control system as in Eq. (6) such that it acts as the identity operation on A i.e. (z) = (ols S2, ,t,a Yo E E Then the system is called E-action-equivarr if
F(o-(s;;;;)) a)) ols F()s, , for (s), E ESXA, VO
(9)
Lemma 3.1 The map 7 Σ x K(S x A) + K(S x A) given by (s+g)(s, - g(0-1- 0. defines a group action on the Koopman space of observables K(S A).
Theorem 3,2 A Koopman operator K. of E-action-equivarian system 81+1 = F(sr, az) satisfies
[o*((g)]($$+ a²) = [K((rRg))(ss a.)
(10)
In particular, it is easy to see that the biliniarisation in Eq. (8) is aaattionequivariiin if fa s) = als fa(s), Vi 0,.-.,m. m. Let thus turn our focus the relevant case of control system 8141 F(sr,a)) az) which admits a Koopman operator description as
g( K((e)g(s1). for a, E A, Vgg E ((S)
(11)
where (K(a)) is family of operators with analytical dependence on A. Note that the bilinearisation in Eq. (8) is a special case of Eq. (11), Furthermore, let (U(a) be a family of invertible operators s.t. U(a)= : K(S) + T(S A) is a mapping to the (Banach) space of eigenfunctions J(S AA) with eigenfunctions (p(s, a) := (a)g(() which obeys (o))(())((()> A(a), with A(a)p(ss a) (a(a))((s a) and where {a : A E. The existence of such operators on an infinite dimensional space requires the Koopman operator in Eq. (11) to be self-adjoint or diite-ddmmesional (approximate) matrix representation to be diagonalizable.
Lemma 3.3 The map Φ Σ x ((S) (U(A))aEA + K(S) given by
(oa - (U (5) g(00
(12)
defines a group action on the Koopman space of observables K(S). Where स is defined analog to Lemma 3.1 but acting on JSS x A) instead of K(S x A) by (a+:)((s, p(o alo--.ss00, s,).. We refer control system in Eq. (11) admitting a d-symmetry as E-action-equivarian..
Theorem 3.4 Let 8141 = F(ss, 0t) be -aaction-equivaiinn control system with symmetry action as in Lemma 3.3 which which furthermore admits a Koopman operator representation as
g( = K((r)g((:)) for a, A,, Vg E ((()
(13)
Then
((C.)))) (s)) [K((at) (g)
(14)
Moreover, a control system obeying equations 13 and 14 is E-action-equivarian locally if g exists for a neighborhood of st. i.e. then o F(s,, a.) = F(a (s,, a.).
'ee Appendix A for the details. We found that, in practice the Koopman operator leaned by the neural nets diagonalizable almost everywhere.
5

---

Under review as a conference paper at ICLR 2022
Let us next provide theoretical statements how data-points may be shifted by symmetry transfor- mations of solutions of ODEs. To establish an easier connection to the next section we introduce the notation Let E: : S + XS) and D K(S) + S denote the C' "diiferrentaable encoder and decoder to and from the finite-dimensional Koopman space approximation, respectively, i.e. FoD = Do = id.
Theorem 3,5 Let $1+1 F(st, ai) be a control system in equation 13 and da, operator obeying equation 14. Then oú. (8, (11,) 1 (5c, St-1; 06) with
s = D((11 + (o..)&(())) 81+1 mannnD(( + 6 < + (Fa.))(((++))
(15)
is one-parameter local Lie symmetry group of ODEs. In other words one can use a symmetry transformation to shift both st as well as 8141 such that $1t1 = F(5,, az).
3.2 KFC ALGORITHM
In the previous section we laid the theoretical foundation for the KFC-algorithm by providing a raod-map how to derive symmetries of dynamical control systems based 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 r, is not part of the symmetry shift process and will just remain unchanged as an assertion.
On the power of using symmetries: Let emphasize the practical advantage of employing symmetries to generate augmented data points. It is evident from Theorem 3.5 that symmetry transformation shifts both s, well 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 < such that Is 8 Oe)) to limit -ittriiiitiiittion generalisation errors Algorithmically the symmetry maps are derived in two distinct ways which we denote KFC and KFC++. The latter, constitute 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. Morcover, in practice one has that a E = id. Thus all assumptions of theorem 3.5 hold approximately.
The Koopman forward model: The KFC algorithm requires pre- trainini of Koopman forward model F S x A- + S which is closely related to a VAE architecture as
D(E(s)) 8t
if 0: VAE
(16)
D (Ko E. K. an.)) F(()) 86+1 if c=== = 1: forward prediction
where both of F. and D are approximated by Multi-Layerr-erren Percepprossssssons (MLP's) and the bilinear Koopman-space operator approximation are implemented by a single fully connected layer for K,,, = 0, o,..., m, respectively. The model is trained on batches of the offline data-set tuples ($1+1, $t, az) 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 al., 2021) our approach leaves the policy improvement Eq. (3) unchanged but modifies the policy optimisation step Eq. (4) as
(++ + Qi (oi, (8e41 aee1) Qi (0 (6rlo))), aa))
)
Q f- arg min Q
E $1,01:111~0
+ +ā D
log Σ exp (Qi(s), acon(s:)) [Q.(s),
(17)
No assumptions the Koopman operator imposed. Moreover, note that the equivalent theorem holds when e Ca. + Ei fi ob, to a local N N-parameter Lie group MMore elaborate studies employing the extended literate on Koopman spectral analysis are desirable. See appendix B for details numerical errors.

---

Under review as a conference paper at ICLR 2022
The state-space symmetry generating function 0² S - S depends on the normally distributed random variables e. 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 level*
KFC oú : 3 = D((1 + coa.) (.a.)((()) E(s))
KFC+ is 8 D oa. (2))
where Jn, (e) Re( (U(e) diag(ei....., U (a:))
(18)
and (a+) diagonalizes the Koopman operator with eigenvalues X; i ,, i... the latter can be expressed as K((+)= = U(sa) /ii.g(A....... \N) (at). Note that we abuse our notation as E(s) = ((() = [g] (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 Cu to commute with the approximate Koopman operator in Eq. (16).
For the KFC option in Eq. (18) the symmetry generator matrix da is obtained by solving the equation Ca, -((a)) - K(aa) Jz = O which may be accomplished by employing Sylvester algorithm (syl). For KFC++ we compute the eigen-vectors vectors of K[(a) which then constitute the columns of Ua).. Thus in particular one infers that [o.,, (e), K((a)]] 0. The latter, is solved by construction of o.(() 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 Px; otherwise use a random normally distributed state shift.
4 EMPIRICAL EVALUATION
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), I). we perform experiments on 6 different challenging robotics tasks from MetaWorld (Yu et al., 2019) and RoboSuite (Zhu al., 2020). We compare KFC 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.
4.1 D4RL BENCHMARKS
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 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 non Mar-Maakovinn behaviour policy (more details can be found in the D4RL manuscript (Fu et al., 2021)). Furthermore. to our knowledge, the results for KFC++ are ofate-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.
'Eq. (18) holds for both 8; as well as 9t+1 thus have deliberately dropped the subscript. Moreover, note that case (II) the symmetry transformations are of the form o,, (é)) = Ei=1 6r oz,-
$1,y 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(é) R-valued matrix.
7

---

Under review as a conference paper at ICLR 2022
Table 1:
We experiment with the full of the D4RL tasks and report the normalized episodic returns over 5 random seeds using the same protocol as et al. (2021). compare against competitive baselines including CQL and the two best performing S4RL-data augmentation strategies. We see that KFC and KFC++ consistently outperforms the baselines. We use the baseline numbers reported in Sinha et al. (2021).
Domain
Task Name
CQL
S4RL-(W)
S4RL-(Adv)
KFC
KFC++
antmaze-umaze
74.0
91.3
94.1
96.9
99.8
antmaze-umazz--iivene
84.0
87.8
88.0
91.2
91.1
antmaze-meciuumpplay
61.2
61.9
61.6
60.0
63.1
AntMaze
antmaaze-mediun-divvrss
53.7
78.1
82.3
87.1
90.5
antmaze-large-play
15.8
24.4
25.1
24.8
25.6
antmaze-large-divers
14.9
27.0
26.2
33.1
34.0
cheetah-randee
35.4
52.3
53.9
48.6
49.2
cheetah-medine
44.4
48.8
48.6
55.9
59.1
Mheetah-medium--rpph
42.0
51.4
51.7
58.1
58.9
cheetah-meddiu-exxie
62.4
79.0
78.1
79.9
79.8
hopper-random
10.8
10.8
10.7
10.4
10.7
hopper-medium
58.0
78.9
81.3
90.6
94.2
Gym
nupper-medium-ii
29.5
35,4
36.8
48.6
49.0
hopper-mediummexpen
111.0
113.5
117.9
121.0
125.5
walker-random
7.0
24.9
25.1
19.1
17.6
walker-medium
79.2
93.6
93.1
102.1
108.0
walker-mediun-reelaa
21.1
30.3
35.0
48.0
46.1
walker-medium-
mepeumeeppprrtt expert
98.7
112.2
107.1
114.0
115.3
pen-human
37.5
44.4
51.2
61.3
60.0
pen-cloned
39.2
57.1
58.2
71.3
68.4
hammer-human
4.4
5.9
6.3
7.0
9.4
hammer-cloned
2.1
2.7
2.9
3.0
4.2
Adroit
door-human
27.0
35.3
44.1
46.1
door-cloned
0.4
2.1
0.8
3.6
5.6
relocate- human
0.2
0.2
0.2
0.2
0.2
relocate-cloned
-0.1
-0.1
-0.1
-0.1
-0.1
kitchen-come sitchen-complete kitchen-paris partial
43.8
77.1
88.1
94.1
94.9
Franka
49.8
74.8
83.6
92.3
95.9
push
pok-laco-ccn
book-1ff
reach wall
80
Parcami of Ccueb Resech
Perdent of Ocaala Reschhdd
MEC+4
COL-94RL
KPC++
COL-64RL
COL-64RL
(a) MetaWorld Environments
(b) RoboSuite Environments
Figure Results on challenging dexterous rohotics environments using data collected using similar strategy as S4RL (Sinha et al., 2021). We report the % of goals that the agent is able to reach during evaluation, where the goal is set by the environments. We see that KFC and KFC++ consistently outperforms both CQL and the two best performing S4RL variants.
4.2 METAWORLD AND ROBOSUITE BENCHMARKS
To further test the ability of KFC, perform additional experiments on challenging robotic tasks. Following (Sinha et al., 2021), we perform additional experiments with 4 World environments (Yu et al., 2019) and 2 RoboSuite environments (Zhu al., 2020). We followed the same method to
8

---

Under review as a conference paper at ICLR 2022
collect the data described in Appendix F of S4RL (Sinha 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 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 guide the data augmentation strategy.
RELATED WORKS
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 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 VAE 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 aate-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 different way. Symmetry-basee 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- et al., 2019) the authors discuss representation learning of symmetries (Higgins al., 2018) allowing for interactions with the environment. A Forward- VAE model which is similar to our Koopman Forwardd VAE model is employed. However, our approach based on theoretical results providing 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 block- diagonal Koopman space representation for non-control dynamical systems is discussed. This is close to the spirit of disentanglement (Higgins et al., 2018). 1)). 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 the first to discuss Koopman latent space for data augmentation.
CONCLUSIONS
In this work we proposed a symmetry-based data augmentation technique derived from a Koopman latent space representation.
enables a meaningful extension of offline RI. 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 Bilincarisation Ansatz for the Koopman operator. However, the empirical results show that the framework is successfully applicable beyond those limitations. We empirically evaluated 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-ae of Q-learning algorithms.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Sylvester algorithm URL
15.1/
htfferenceeggnerattee///// wtw.ddp.....l....oleeeesseee// solve_sylvesteer html.
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning, 2020a.
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. optimistic perspective on offline reinforcement learning. In International Conference Machine Learning, 2020b.
Ravindran Balaraman and G. Barto Andrew. Approximate homomorphisms: framework for non- exact minimization in markov decision processes. In In International Conference on Knowledge Based Computer Systems, 2004.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym, 2016.
Steven L. Brunton, Marko Budišić, Eurika Kaiser, and J. Nathan Kutz. Modern koopman theory for dynamical systems, 2021.
Hugo Caselles-Dupré, Michael Garcia Ortiz, and David Filliat. Symmetry-based disentangled representation learning requires interaction with environments. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6( 6(18):503-556, 2005.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Proceedings of the 36th International Conference on Machine Learning, volume 97
of Proceedings of Machine Learning Research, pp. 2052-20020000 2062. PMLR, 09 15 Jun 2019.
Dimitrios Giannakis. Data-driven spectral decomposition and forecasting of ergodic dynamical systems. Applied and Computatiional Harmonic Analysis, 47(2):338-333 396, 2019. ISSN 1063- 5203. doi: http:://oi..ry//0000/////aaaaaat fam//scienee
URL https: /www. sciencedirect.
Debdipta Goswami and Derek A. Paley. Global bilinearization and controllability of control-affine nonlinear systems: A koopman spectral approach. 2017 IEEE 56th Annual Conference on Decision and Control (CDC), pp. 6107-61 12, 2017.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 31, Curran Asso- ciates, Inc., 2018. URL httpt:tttt//ttt ftctps://prooeeedings new.eeeddddnnnnnnnnneeeeeeeeeeee http:/ssrsssp....ppppeee///////e cccppaper/2018/Eleee 2de5d1 (ec.0a66226600066222222222222000 Papere pdf.
Tuomas Haarnoja, Haoran Tang. Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy- based policies. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1352-1361. PMLR, 06-11 Aug 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1861-1870. PMLR, 10-15 Jul 2018.
Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loie Matthey, Danilo Rezende, and Alexander Lerchner. Towards definition of disentangled representations, 2018.
10

---

Under review as a conference paper at ICLR 2022
Eurika Kaiser, J. Nathan Kutz, and Steven L. Brunton. Data-driven discovery of Koopman eigenfunc- tions for control. Machine Learning: Science and Technology, 2:035023, 2021.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, and Sergey Levine. Scalable deep reinforcement learning for vision-based robotic manipulation. In Aude Billard, Anca Dragan, Jan Peters, and Jun Morimoto (eds.), Proceedings of The 2nd Conference on Robot Learning, volume 87 of Proceedings of Machine Learning Research, pp. 651-673. PMLR, 29-31 Oct 2018.
Diederik P Kingma and Max Welling. 131iv:1312.611 2013.
Auto-encoding variational bayes.
arXiv preprint
Bernard Osgood Koopman. Hamiltonian systems and transformation in Hilbert space. Proceedings of the National Academy of Sciences of the United States of America, 17(5)3115318, 1931.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1179-1191. Curran Associates, Inc., 2020.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representa- tions for reinforcement learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 5639-5650. PMLR, 13-18 Jul 2020a.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Re- inforcement learning with augmented data. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 19884-19895. Curran Associates, Inc., 2020b.
Lihong Li, Wei Chu, John Langford, and Robert Schapire. A contextual-bandit approach to personalized news article recommendation. In Proceedings the 19th international conference on World wide web, pp. 661-670, 2010.
Yunzhu Li, Hao He, Jiajun Wu, Dina Katabi, and Antonio Torralba. Learning compositional Koopman operators for model-based control. In Proc. of the 8th Int'l Conf. on Learning Representation (ICLR'20) 2020.
Alexandre Mauroy, Igor Mezić, and Yoshihiko Susuki. The Koopman Operator in Systems and Control: Concepts, Methodologies, and Applications. Springer, 2020.
I. Mezić. Spectral properties of dynamical systems, model reduction and decompositions. Nonlinear Dynamics, 41: 1::0-3335, 2005.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In leml, 2010.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32: 8026-8037, 2019.
Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. arXiv preprint arXiv: 709. 10087, 2017.
Anastasiya Salova, Jeffrey Emenheiser, Adam Rupe, James P. Crutchfield, and Raissa M. D'Suzz.. Koopman operator and its approximations for systems with symmetries. Chaos: An Interdisci- plinary Journal of Nonlinear Science, 29(9):093128, 2019. doi: 10. 509999.
Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4RL: Surprisingly Simple Self- celf-Supervision for Offline Reinforcement Learning. In Conference Robot Learning, 2021.
11

---

Under review as a conference paper at ICLR 2022
Subhrajit Sinha, Sai P. Nandanoori, and Enoch Yeung. Koopman operator methods for global phase space exploration of equivariant dynamical systems, 2020.
Elise van der Pol, Daniel Worrall, Herke van Hoof, Frans Olichoek, and Max Welling, Mdp homo- morphic networks: Group symmetries reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 4199-4210. Curran Associates, Inc., 2020.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. 2019.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. arXiv preprint arXiv: 1910:1910010097, 2019.
Yuke Zhu, Josiah Wong, Ajay Mandlekar, and Roberto Martin-Martin. robosuite: A modular simulation framework and benchmark for robot learning, arXiv preprint arXiv:2009.12993, 2020.
12

---

Under review as a conference paper at ICLR 2022
A KOOPMAN THEORY AND SYMMETRIES
This section provides some additional introductory material on Koopman theory and symmetries of dynamical systems.
Lemma A.1 The map x K(S) - K(S) given by (a g)(s) + g(a=l. defines a group
action on the Koopman space of observables K(S).
Theorem A.2 Let K. be the Koopman operator associated with E-equivariant system 8141 Then
F(r)).
(o * (s) [K[o * *s)
(19)
Theorem A.2 states that for E-qquivariant 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 (a))AAe be a family of invertible operators s.t. U(a): K(S) + F(S x A) mapping the (Banach) space of eigenfunctions g(s. a) == (s))((s) E F(S x A). Which moreover obeys ((a) A A(a), with A(a)p(ss, a) (a) a) and where (a) AA R. The existence of such operators puts further restriction on the Koopman operator in Eq. (11). However, as our algorithm employs the mnite-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-dimensionall case we employ an alternative formulation of the so called spectral theorem below.
Theorem A.3 (Spectral Theorem) Let K. be bounded self-adjoint operator a Hilbert space H. Then there is a measure space (S, Σ, (1) and a real-valued essentially bounded measurable functioo on S and unitary operator U H. (S) i.e. U = U id such that
UAU* K. with ({))ss)
(20)
In other words every bounded self-adjoint operator is unitarily equivalent to a multiplication operator. In contrast to the dinee--immnsionll case need to slightly alter criteria to (aa))(a))(()=++ A A(a), with A(a (s)) a) = (a, s) a) and where (s a) S x A + R. Concludingly, a sufficient condition for our criteria to hold in terms of operators Hilbert spaces is that the Koopman
operator K((a) is self-adjoint i.e. that
K(a) K((a)*
(21)
B PROOFS
In this section we provide the proofs of the theoretical results of Section 3.
Proof of Lemma A.I and Theorem A.2 The proofs of the lemma as well as the theorem can be found in (Sinha al., 2020).
Proof of Lemma 3.1: We aim to show that map म Σ x K(S x A) - K(S x A) given by (oig)(s, a) + g a) defines group action the Koopman space of observables K(S A), where Σ defines the symmetry group of definition 5. Firstly, let g E K(S x A) and op E Σ be the identity element, then we see that it provides existence of identity element of by
(oozg) a) = g(ap" s, a) = g(s,a)
Secondly, let 01, 02 E E and @ denoting the group operation i.e. on @ 02 03 E Σ. Then
(02 (s) = σ. (g(0T) s.a) = 9 (oi s). g o ) 5),
g((و -1 a) ((0) v وو)=مg (8, a)
و3مم (s,a)
where in (I) we have used the invertibility of the group property of Σ. Lastly it follows analogously that for o, o E that
( )s.)))ss.))) (s, = (a:))=((ooooo ((- = (0otg) (s, g(s.a)
Thus the existence of an inverse is established which concludes to show the group property of s.
13

---

Under review as a conference paper at ICLR 2022
Proof of Theorem 3.2: We aim to show that with K. being the Koopman operator associated with E-action-eed nnetion-equivariant system $141 = ((st,aa)) Then
(KCg) (s, a) = [ (s.a)
First of all note that by the definition of the Koopman operator of non-affine control systems obeys
[Kg](s;; ((F(ss;: ar-1)
Using the latter one thus infers that
(Kg) (st, a+) (Fg(F((s, 0z), ac+1) -1----1----- F(s,, ar) a(11 _ g(F(o St, at), a+t1);
where in (I) we have used that it a E-action-equivvaiint system. Moreover, one finds that
F(F(o 8t, az), 1++1) Kylo -tor-t1._ 1., at) (((ogg) (st,
which concludes the proof.
Proof Lemma 3.3: We aim to show that the map á Ex xx()) x {U(a))aEA + X(S) given by
(a (+)
( - (1 (((((((()))))) (0)=
defines a group action on the Koopman space of observables K(S). Where * is defined analog to Lemma 3.1 but acting on I(S A) instead of K(S A) by (onp)(s, - s(a 1. s,a). First of all note that
(a) (U(a))))) (a) p(s. U- (a)p(oo s, g
1 s) (22)
We proceed analogously in the proof of Lemma 3.1 above. Firstly, let E K(S) and dg E Σ be the identity element then on infers from Eq. (22) that it provides existence of an identity element by
(Fa.o&g)(s) = g0.5 .)) = g(s)
where we have used the notation
(o.,) = (0- (s), for i == 0,...
Secondly, let 01,02 E E and denoting the group operation i.c. or o 02 03 E Σ. Then
= = (ouzt(0a1)0)) (s)
s) -(05' .o 5)) = 9((0z" o ooT s)
@ o U 02) 8) g 8) ( (s)
where in (I) we have used the invertibility of the group property of >. Lastly, it follows analogously that for a, o E Σ that
(o = (o)))))) = (ookg) = g(s)
Thus the existence of an inverse is established which concludes to show the group property of d.
Proof of Theorem 3.4: We aim to show twofold.
Firstly, that with 8411 = (st, 0z) be a -action-equivariian control system with a symmetry action in Lemma 3.3 which furthermore admits Koopman operator representation as =:
g(s++1)= = K((x)g(ss) for a, E A, Aig E K(S)
Then
[aa, (KC(a)))) (sr) [x((a)) (oa, kg) (81)
14

---

Under review as a conference paper at ICLR 2022
Secondly, the converse. Namely, if control system 81+1 F(st. at) obeys Eqs. (13) and (14), then it is E--action-equivarian i.e. F(t,, az) = F( (st, a,))- For notational simplicity drop the subscripts referring to time in the remainder of this proof i.c + s and a₁ - 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))](s) (F((s)))
Using the latter one infers that
aa a (IC(a)))) (s)
a,
(0(00 ((s.a))))))
s(F(o-'s,a))
Moreover, one derives that
( s, {(a)g(o-1s)
(a)
= (a)
= ((a)(
(o)
[K(a) (a (o.*g)
which concludes the proof of the first part of the theorem. Let us next show the converse implication the discrete time-system notation explicitly. Let Σ be
i.e. < For this case it is practical a symmetry of the state-space and be š; = 5 and š = o St+1 the o-shifted states. Then
g (3c+1)
g ((((Fe)))))
K((a)) g(s))
K((a))((o s.)
= K(aa)) (1la))
[((a))(a0*1) (8x)
i ( (3x)
Thus in particular
0a =(gg(8041)
0
1 (0 (K((x))))
(3,)
Lemaa 3.3
((ax))) (5x)
15

---

Under review as a conference paper at ICLR 2022
where in (I) we have used that the symmetry operator commutes with the Koopman operator. More- over, one finds that
Ja (+41)) - s(+1) g(0 8(41) g((+11)
from which one concludes that
g( = [K(ar)g] (5t) ) = g (3 a:))
Finally, we use the invertibility of the Koopman space observables i.e. g (g())) = 8 to infer
Be+1
F(5., a4)
However, in general s the Koopman space observables are not invertible globally. The "inverse function theorem" guarantees the existence of local inverse if g(s) is C³ differentiable for maps between manifolds of equal dimensions. However, we assume inevitability locally.
S--1 = J St+1
F(5,, ae) = = F(st, ae) = F(a St, oz) ) = o() (23)
which at last concludes our proof of the second part of the theorem.
may account for practical limitations i.e. error by the assumption (o K[(a)]] ea 1. One then finds that $ = F(,, ae) + OEa.). Thus the error becomes suppressed when Ea < e. The error may be due to practical limitations of capturing the true dynamics as well as the symmetry map.
Extension of Theorem 3.5: Moreover,
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(sr, at) be vuction-qquivariant control system in Theorem 3.4. For symmetry maps
3 = D + co.,) $111 = =D((1 + +co..)&E(s4t1)
(24)
we aim to show that one can use symmetry transformation shift both 8t as well 8t+1
oy (st, $141,a;)
(5 $141, at) s.t. 811 = a()
(25)
From equation 24 and the definition of the Koopman operator one infers that
h1+ = D oa, E((++1) = D( (,,))(()))(ss)))
(26)
By using that [aa., K(aa)) = O and equation 24 one finds that
5++1= D K(ar)o,, #E(ss) = D( K((a)) E(D(o,,$E(s))) E((s)))) D( ((aa))(3)))
(27)
which concludes the proof of the first part of the theorem with da, + (1 + co..).
Local diffeomorphism. Per definition the maps D, E' are differentiable and invertible which implies that they provide local diffeomorphism from and to coopman-space Also the linear map i.e. matrix multiplication by (1 + co.,) a diffeomorphism. Thus the symmetry map of i.e. Eq. (15) constitutes local diffeomorphism. The above proof culminating in Eq. (24) implies that we have local diffeomorphism mapping solutions of the ODEs to solutions, thus a symmetry of system of ODEs.
Limitations. However D, E to be invertible a strong assumption as it implies that the Koopman space approximation admit the same dimension the state space. Note that however in practice as we only require D zs id one may choose other Koopman space dimensions. Local Lie group. What is left show is that the symmetry of the system of ODEs locally is Lie group. In definition 2 we need to show points (1)-(3).
For E = one finds that of is the identity map i.c for such that
(5( 0) = D(11 + 0 oa.) = D(E(s:))
(28)
For notational simplicity study the general case + E.ai) Gat-
16

---

Under review as a conference paper at ICLR 2022
2.
00 = +(1 (20,,)((s)
(29)
= D((1 + + + + O(a 62)
= D((1+ + +11 t2 + (2)
= outo
for every leil: |eal <
3. Se(s,, c) admits Taylor series expansion in c, i.e. in neighbourhood of determined by
E = O. We may Taylor expand D around the point L(ss) = as
N
3D
8e(8t,)) = D(E(st)) +
aa, E(sr)) Ogr
+ O
(30)
1=1
thus
5e(3,, = + e{(st) - O(E2)
(31)
for (($t) = (o,, Ei(s))) Dgr
4. The existence of inverse element (o )-_))-1-0 oe 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 oo, K((a)) = one finds $c+1
F(5r, at) - Olea). Note that Ca, + (1 + E ET..) admits an expansion the parameter e. For << one can Taylor expand the differentiable functions D, E to find that
OD gi
8t = 8t + 8(c), with 8(e) = Σ
Ögi B(s.) 0(22)
(32)
where
N. The analog expression holds for t + 1 i.e. 5+41 = St+1 + Oee). af ar(st) is differentiable function of the states? one may Taylor expand the Koopman operator using equation 32 as
and g Ess). and with the index / = ,.
((a(()))) = ((a{(s))) + ΣΣ 5n(() ah (8t) KG = K(a{(s)) + +cA+
(33)
n
where and (8))
ls, and I,..., dim(sE) and with
aD,
A = ΣΣΣ>>>
9i
Ogr B(se) d'n Ks
(34)
However, by the definition of the symmetry map (3))) = at (st), thus the Koopman operators in equation 33 match Vc. By using the assumption on the error infers from equation 26 that
š141 ((11 co.,, (((())))s)))
D((K(a))+
ca ca (1 ca )+E(s))
((K(O)) (1, + E(st) + O(Eae)
F(K))) +. F(ss)))) + O(eac))
(35)
One may Taylor expand D by making of 6a aas
0D 0gi ((a.)
$1+1 = D(((a))(()) + ca EE((s))
+ +0(cct)+0(e))
(36)
Although neural networks modeling the policies modeling the data-distribution contain dinfeffeereniabeee components they differentiable almost everywhere.
'The reader may wonder about the explicit use of the index é, (c) expressing the dimension of the state. it is necessary here although we have simply used s, implicit vector without explicitly specifying any components throughout the work.
17

---

Under review as a conference paper at ICLR 2022
Thus under the assumption that the numerical error is given by violation of the commutation relation one finds 6(+1 F(,,, at) + O(ea) where
OD
Σ Er Br( Ögr K()))(())
~ O
(37)
By comparison with equation 32 one infers that the e-expansion a good approximation to the real dynamic if 6a < This concludes the proof of theorem 3.5. Let note that the proof for state shifts arising from a sum of symmetry transformations as
N
N
8 D(1
Σ ((od)))((s))))) 1=1
6+1=D( + [eIod))E((s4t)))) 1=1
(38)
is completely analogous, and an analogous theorem holds. We choose the number of symmetries to 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 Koopman forward model F F:SS A + as
D(E(si)) 8t
if 0 corresponds VAE.
(s, ax)
D (K0 + Kr (aa())(s)) 8
if c 1: forward prediction model.
(39)
where both of E and D approximated by Setti-Layer-Percettron (MLP's) and the bilinear Koopman-space operator approximation are implemented by a single fully connected layer for K===0..... respectively. The model is trained on batches of the offline data-set tuples (8,41;8;,a) and optimized via the loss-function
T2(F 0z), 8141) + 72 12 + at), 8t 1+.+0 +
(40)
where 8, is normally distributed random sate vector shift with zero mean and standard deviation of 6- 10-2 where the lo-so is the Huberloss. For training the Koopman forward model in Eq. (16) we split the dataset loop over the dataset to generate the required symmetry transformations and extend the replay buffer dataset with the latter according Eq. (18) as
randomly selected training and validation set with ratios 70%/30%. We then
KFC:
(st, (e:11,0:71
symmetries symmetries
(s 8141, at, Tt, Go.)
(41)
KFC++:
(St, St41, Qc:
(st, a,, Tr, U Uat))
is more economic to compute the symmetry transformation pre-processing step rather than runtime of the RL algorithm.
Let us reemphasize couple of points addressed already in the main text make this section self- contained. Note that we have built upon the implementation of CQL (Kumar et al., 2020), which based on SAC (Haarnoja al., 2018). We also use the same hyperparameters for all the experiments presented in the COL and S4RL papers for baselines.
Koopman forward VAE-model: For the encoder as well as the decoder we choose a three layer MI.P architecture, with codden-dimensions. (512, 512, N) and [512, 512, m.) respectively. Where N :2 is the Koopman latent space dimension and m the state space dimensions following our notation from section 2. The activation functions ReLU (Nair & Hinton, 2010). We use ADAM optimizer with learning rate 3- 10 We train the model for 75 epochs. The batch size is chosen to be 256. Lastly, we employ random normally distributed ttateaaugmentttion for the VAE training i.e. s + $ + Se with s, E EN(01 1. 10-2).
Hy Hyper-parametert et al., 2020) except for the fact that we do not use automatic entropy tuning of the policy optimisation step Eq. (3) but instead fixed value of 0.2. The remaining hyper-paramettrss of the conservative
KFC: We take over the hyper-parameee settings from the papea (Kumar
18

---

Under review as a conference paper at ICLR 2022
Q-learning algorithm are follows: y 0.99 and == = 10-3 for the discount factor and target smoothing coefficient, respectively. Moreover, the policy learning rate 1 10 and the value function learning rate is 3 10 for the ADAM optimizers. We enable Lagrange training for å with threshold of 10.0 and learning rate of 3 10-4 where the optimisation done by an ADAM optimizer. The minimum Q-weight is set 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 the value function are three layer MLP's with hidden dimension 256 and DeeL--activatton. The batch size is chosen to be 256. Moreover, the algorithm performs behavioral cloning for the first 40k training-stege steps, i.e. time-steps in an online RL notation.
The KFC specific choices are as follows: the split of random to Koopman-symman symmerry state shifts 20/80 i.e. Pk 80%, whereas the option for the symmetry-type is either "Bigenspace" or "Sylvester" referring to case (I) and (II) of Eq. (18), respectively. Regarding the random variables in the symmetry transformation generation process, we choose & E (O, 10-4) for case (II) and E (0, 5 10-5) for case (I) in Eq. (18) after normalizing by its mean. While for the random shift we use E (0 3 10 5) which the nyper-parameter used in (Sinha et al., 2021).
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) Xcon(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.
Computational setup: We performed the empirical experiments
D ABLATION STUDY
Task Name
S4RL-(N)
KFC
KFC+
KFC++-contact
cheetah-randus cheetah-medium
52.3
48.6
49.2
49.4
48.8
55.9
59.1
61.4
cheetab-meddiu--eella
51.4
58.1
58.9
59.3
cheetah-meddun-expper
79.0
79.9
79.8
79.8
hopper-random
10.8
10.4
10.7
10.7
hopper-medium
78.9
90.6
94.2
95.0
hopper-medium--eplayy
35.4
48.6
49.0
49.1
113.5
121.0
125.5
129
walker-random
24.9
19.1
17.6
18.3
valker-medium
93.6
102.1
108.0
108.3
walker-mediune medium---mdi
30.3
48.0
46.1
48.5
malker-medium-experr
112.2
114.0
115
118.1
Table 2: We study the effect of combining S4RL-V based augmentation training on "contact" event transitions (state transitions where the agent makes contact with surface, read D.1 for more details), and KFC++ aug- mentation training on non-"con "contact" events. We CFC++contact performs similarly slightly hetter than KFC++ baseline, which is expected. We experiment with the Open Al gym subset of DARL tasks and report the mean normalized returns over 5 random seeds.
most cases,
D.1 TRAINING ON NON CONTACT EVENTS
To test the theoretical limitations in a practical setting one interesting experiment use the S4RL- N augmentation scheme when the current state is "contact" state and to choose KFC++ otherwise. We define a "contact" state as one where the agent makes contact with another surface; in practice looked the state dimensions which characterize the velocity of the agent parts, and if there is sign difference between š, and $t41, then contact with surface must have been made. The details regarding which dimension maps agent limb velocities is available in the official implementation of the Open AI Gym suite (Brockman et al., 2016).
Note that the theoretical limitations twofold, firstly the theorems 3.4 and 3.5 only hold for dynamical systems with differentiable state-trae nnte-transitions; secondly, we employ a Bilincarisation Ansatz for the Koopman operator. Contact events lead to non-continuous state transitions and moreover
19

---

Under review as a conference paper at ICLR 2022
would require incorporating for external forces in Koopman description, thus both theoretical assumptions are practically violated in our main empirical evaluations KFC and KFC++ By performing this ablation study which is more suitable for our Ansatz, i.c the KFC++ part is only trained on "differentiable" trajectories without contact events. In other words, the Koopman description is unlikely give good information regarding 8++1 if there is "contact event", 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 OpenAl Gym subset of the D4RL benchmark test suite. We denote the proposed scheme KFC++ccontact and compare both S4RL-N and KFC++ We see that KFC++ccontat performs comparably and slightly better in some instances, compared to KFC++. This 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++-contac. However, we do not say this conclusively.
D.2 MAGNITUDE OF KFC AUGMENTATION
To further try to understand the effect of using KFC for guiding data augmentation, we measure the mean 1.2-norm of the strength of the augmentation during training. We measured the mean 1.2normm between the original state and the augmented state, and found the distance to be 2.6 x 10-4 For comparison, S4RL uses fixed magnitude of 1 10-4 (Sinha 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 hyperparameter search decision.
D.3 SYMMETRIES
In this section we provide 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
(42) averaged over all sampled data points (D4RL) is comparable between the two setups. Moreover simply use the encoder We then compare how well the shifted states are in alignment with the actual dynamic using the online Mujoco environment To do so set the environment state to the value 8, and use the action a, to perform the environment step which we denote by 81+1 = (st, 0z) The performance metric is then given by AE == ===11:111 M(s, a+))
AS = 55 sil + 18 st+1|
(43)
Note that AE is simply the error to the actual dynamic the environment. We take the model trained on the medeee--mdiieeerrr 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 1 (3,, a,) 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 qualitative difference between the distributions obtained. Most notable the symmetry shift allows for a controlled shift by a larger AS while still maintaining a small error AE This amounts 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.
D.4 PREDICTION MODEL
As a further ablation study it is of interest to compare the symmetry induced state shifts tuple
(ot, 3141) to one obtained by the forward prediction of our VAE-forward model equation 16.
9 Note that this requires a minor modification of the environment code by a function which allows set states to specific values.
M (St, ax) denotes the step done by the active gym environment.
20

---

Under review as a conference paper at ICLR 2022
random shift
symmetry shift
175
150
03
125
1.00
0.75
0.50
025
0.0
0.2
0.4
0.6
08
10
0.00
0.0
02
0.4
0.6
10
150
125
1.00
075
025
10
0.0
(a) Hopper position state shifts
(b) Hopper velocity state shifts
Figure 3: Symmetry shifts vs. random latent space shifts of same magnitude in the mean of AS compared by its accuracy in the online evaluation in the gym Hopper-v3 environment. The first line shows the result where st and 8++1 are shifted by the same N-dimensional random variahle. For the ablation study in line twice from random variable
sample
shifts 8; 8t+1, respectively.
KFC+tpredictio: In particular use the forward prediction on 5t obtained by the KFC++ shift, ie.. 5e+1 = (3r (3r, We refer to this setup in the following as KFC++prediction. See table for the results.
Fwd-prediction: Moreover, we study the case where we š, is obtained by shift with a normally distributed random variable (0,0 6 10 10-3) and then simply forward predict as š;+1 F (3,, az). This is comparable to conventional model based approaches employing our simple VAE forward model. This constitutes ideal systematic comparison 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 distribution values.
More interesting however is that Fwd-prediction falls behind both KFC as well as KFC++ as well KFC++prediction. This is strong indicator that the symmetry shifts provide superior out-of-distribution data for training a Q-learning algorithm.
D.5 DISCUSSION
Although our KFC framework is suited best for the description 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 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 general non-affine control system admits bilinearisation (Brunton et al., 2021)
21

---

Under review as a conference paper at ICLR 2022
Task Name
KFC
KFC++
KFC++-prediction
cheetah-random
48.6
49.2
46.5
cheetah--ediium
55.9
59.1
53.7
chectab-meddim--eplly
58.1
58.9
55,3
Sreetah-medium-expee
79.9
79.8
76.3
hopper-random
10.4
10.7
10.8
hopper-medium
90.6
94.2
90.5
hopper-meddium-eelly
48.6
49.0
44.2
hopper--medium-exppe
121.0
125.5
121.2
walker-random
19.1
17.6
15.6
walker-medium
102.1
108.0
105.3
mellurr-mediuumrepllay
48.0
46.1
45.2
walker--medium-exper
114.0
115.6
114.5
Table 3: Results with prediction model prr+iipciiittion on the Open Al Gym subset of the D4RL tasks, report the mean normalized episodic rewards over random seeds similar to the original D4RL paper Fu et al. (2021).
Domain
Task Name
KFC
KFC+
Fwd-prediction
antmaze-umaze
96.9
99.8
92.7
antmare-umare-diivrse
91.2
91.1
90.1
antmaze-medium-phae
60.0
63.1
60.8
AntMaze
autmaze-medium-diives.
87.1
90.5
88.0
antmaze-large-play
24.8
25.6
23.1
antmaze-large-dives
33.1
34.0
29.3
cheetah-random
48.6
49.2
50
cheetah-medium
55.9
59.1
50.1
Mheetab-mediimm-repee
58.1
58.9
56.4
cheetab-mmedium-exper
79.9
79.8
71.3
hopper-random
10.4
10.7
10.4
hopper-meding
90.6
94.2
82.3
Gym
hopper-mediumt
replayy
48.6
49.0
40.8
121.0
125.5
120.3
walker-random
19.1
17.6
18.4
walker-medium
102.1
108.0
103.2
walker-mediuum-replay
48.0
46.1
41.7
walker-meddium-exper
114.0
115.3
111.8
pen-human
61.3
60.0
49.4
pen-cloned
71.3
68.4
50.2
hammer-human
7.0
9.4
6.1
hammer-cloned
3.0
4.2
4.2
Adroit
door-human
44.1
46.1
41.8
door-cloned
3.6
5.6
1.2
relocate-hume
0.2
0.2
0.2
relocate-cloned
-0.1
-0.1
-0.1
kitchen-complett
94.1
94.9
90.0
Franka
kitchen-partin
92.3
95.9
84.6
Table 4: Results with prediction model Fwd- predictiio the D4RL tasks. We report the mean normalized episodic rewards 5 random seeds similar to the original D4RL paper Fu et al. (2021).
it generically requires the observables to depend on the action-space variables implicitly. Secondly, the Koopman operator formalism is theoretically defined by its action infinite dimensional
22

---

Under review as a conference paper at ICLR 2022
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 procpeesssing which is computationally expensive, i.c. solving the Sylvester or eigenvalue problem for every data-point. Moreover, the Koopman forward model in Eq. (16) serves as function approximation two distinct tasks. Thus one faces 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 stress crucial point. Algorithmically the symmetry maps are derived in two distinct ways KFC and KFC++. The latter, constitute simple starting point 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 more complex latent space descriptions such as e.g. world models. Both on theoretical 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 gains of Q-learning algorithms. Note that currently our VAE model is of very simple nature and the symmetries are extracted in 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.
well their induced performance
23


Reviews:
[{"review_id": "wYExqi1gbX", "reviewer": "Reviewer_ohJ3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\n### Strengths\n- The idea of using learned symmetries of the environment transition operator to do data augmentation in latent space is an intriguing one and to my knowledge is novel.\n\n- The work provides sufficient background that a reader with a background in RL and with a passing familiarity with linear dynamical systems could understand the methodology and goal of the paper.\n\n- I didn’t notice any obvious mistakes in the proofs of the main theorems, though there were several typos (e.g. $ \\bar{\\sigma} *$ in place of $\\sigma \\bar{*}$).\n\n- Empirically, the method does seem to outperform the baselines the authors chose.\n\nWeaknesses:\n- The paper suffers from **lack of clarity**: some symbols are not defined and the types of certain functions and operators must be inferred by referring back to their initial definition. For example, the notation U(a)g used in equation 13 is ambiguous at first as to whether it is meant to refer to a composition of operators or to multiplication of the outputs.\n    - In lemma 3.3, $\\bar{\\Sigma}$ is not defined — further, using the most recent definition of the group would lead me to believe it acts on state-action pairs, however in Lemma 3.3 it is applied to individual states.\n    - The paper states that epsilon is sampled from normal distribution, but doesn’t state the precise variance, only claiming epsilon <<1.\n    - The operator $\\sigma^\\epsilon_{a_t}$ is not defined in Theorem 3.7, nor is it clear how it relates to the (presumed) group element $\\sigma$.\n    - The notation $\\mathcal{F}(S \\times A)$ is not defined. Further, the type/definition of U(a) doesn’t make sense to me: what is the value of (U(a)(g))(s’, a’) ?\n    - the notation $\\bar{\\Sigma}$, $\\hat{\\Sigma}$, etc. should be formally defined. Further, in many definitions the notations $\\phi$ and $*$ are used interchangeably. \n\t\t\n\n- There is a **large gap between the assumptions used in the theoretical results and the empirical setting**, both in terms of the assumptions required and in the outcome being measured. The theoretical results concern the equivariance properties of the Koopman operator, while the empirical results concern the performance of a learned model. \n    - The proposed method looks suspiciously similar to a VAE objective where the perturbation is defined by sigma. Thus it is not clear to me whether the properties of the learned transition model + invariances are actually important to the proposed success of the method, or indeed whether the learned symmetries can really be said to be true symmetries of the Koopman operator as opposed to simply latent-space perturbations\n    - The existence of the family of operators U(a) used in Lemma 3.5 is not shown in main paper, and indeed additional conditions are discussed in the appendix in order for these operators to exist. It is not shown that these conditions are satisfied in the settings the paper is interested in.\n    - It is not clear what types of symmetries exist in the evaluation environments, or whether KFC is able to learn a model that is equivariant to them. \n    - The jump from the theoretical results on equivariance to the empirical evaluation on performance is quite abrupt and it is not clear how the theoretical results are meant to indicate that the method will indeed improve performance.\n    - It’s not clear how the learned encoder/decoder/observables used in practice satisfy the assumptions laid out in the theoretical section.\n\n\n- **The empirical results are not convincing.** The observation that S4RL outperforms KFC on the -random split where model-learning is less useful suggests that most of the benefit of KFC is coming from the model learning component of the algorithm. S4RL and CQL are both model-free methods, so this seems to suggest to me that it might not be the learned latent-space symmetries but rather the model-learning auxiliary task that is making KFC outperform the selected baselines. Additionally,  a number of algorithmic details are missing from the description of the method, making it difficult to interpret the experimental results. \n    - I couldn’t find a description of KFC++ anywhere in the paper.\n    - How is the matrix representation of the koopman operator obtained?\n    - How many training steps do the methods get? Does the model-training phase count towards the  KFC method’s “optimization budget”?\n\n### Potential Improvements\n\n\n- Evaluation against stronger baselines that incorporate model learning and equivariance in the latent space (e.g. DeepMDP, DBC).\n\n- Rewriting the paper to improve the clarity: making sure the types of functions and operators are always clear, using consistent notation, explicitly stating the necessary assumptions for the theoretical results to hold, etc.\n\n- Using more targeted experiments to study the learned symmetries of the model will provide greater insight into whether the theoretical results hold in the latent space of KFC agents.\n\n\n### Notes\n-  I didn’t understand the claim “an exploration of the environment’s phase space” (section 1, page 2)\n- In environments that are not equivariant to any symmetry group, how would you expect the method to behave? \n- There are a number of typos in the manuscript: “the existence on nature” on p2, ’raod-map’ on p6, “regulizer” on pp 3 & 6, “in a the free” just after eq 8, etc. These were separate from my concerns about clarity.", "summary_of_the_paper": "This paper aims to improve generalization and data efficiency in offline RL using a latent-space data augmentation approach inspired by results from Koopman operator theory. The resulting method is evaluated on several environments against a similar method which uses a form of self supervision via data augmentation.", "summary_of_the_review": "Overall, my concerns over the clarity of the paper, the gap between theory and experimental results, and the limited comparison to existing methods lead me to recommend that the paper be rejected.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "XaXBL1oGz6L", "reviewer": "Reviewer_SsNL", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "-The general idea and the results of this work are interesting and seem to be novel. \n\n-However, I found the paper to be rather hard to understand. Some important assumptions are not so clear and require more clarification. Also, some points and notations should be better explained.\n\n\n**The main issues that, in my view, could improve the paper are:**\n\nFirst, there are some grammatical mistakes as well as missing words and typos, for instance\n- Page 5, line 23 and page 13, line 29:    .Which ---> which \n- On page 4, line 9 in the proof of Theorem 3.6, the sentence \"First of all note that by the definition the Koopman operator it obeys\" is grammatically incorrect\n- Page 16, line 1:  us ---> use\n- Page 16, the second line below eq. (27): an controllable ---> a controllable\n\nSecond, some important assumptions are not so clear and require more clarification:\n- On page 4, line 6, it is not clear what are f_i functions. Also, considering  \"f(s,a)=f_0(s) + \\Sum f_i (s) a_i \", it is not clear how eq. (7) can be written as eq. (8). Moreover, writing \"f_i=0, ..., m\" is a bit confusing.\n\nThere are also some comments on Theorems and their proofs:\n- In Theorem 3.4, eq. (12), the first term is \"g(s_{t+1})\" or \"g(s_{t+1}, a_{t+1})\"? Also, in this theorem it is not clear under which conditions we can obtain eq. (12). So, more clarification in this regard is needed (this is also needed for eq. (14) in Theorem 3.6).\n- In the proof of Lemma 3.3, only  \"σ_0∈ ∑\" is the identity element but not \"g\", although it is written \"Let σ_0 and g be the identity element\". This should also be considered in the proof of Lemma 3.5, i.e. only  \"σ_0∈ \\bar{∑}\" is the identity element. \n- In the proof of Theorem 3.4, line 7, \"σ^\" should be changed into \"σ\".\n- In the proof of Theorem 3.6, line 4, what is \"a\"? Is it \"a_t\"? Also the first term is \"g(s_{t+1})\" or \"g(s_{t+1}, a_{t+1})\"?\n- In the proof of Theorem 3.7, line 5, \"σ_{ε}^{a_t}\" should be better explained and needs to be clarified. \n", "summary_of_the_paper": "The paper proposes a symmetry-based data augmentation technique derived from a Koopman latent space representation based on some theoretical results on symmetries of dynamical control systems and symmetry shifts of data. The authors empirically evaluated their method on several benchmark offline reinforcement learning tasks D4RL, Metaworld and Robosuite and showed that their framework consistently improves the state-of-the-art of Q-learning algorithms.", "summary_of_the_review": "This paper needs to be written in a more precise way. Especially, some important assumptions are not so clear and require more clarification.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "NqCCHUz_LN7", "reviewer": "Reviewer_VgB7", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. The paper provides a principled way to generate data augmentation for offline RL (and for dynamical systems in general)\n2. The paper improves over past data augmentation methods (S4RL variants) and show improved performance on D4RL tasks, meta world and robosuite environments\n\nConcerns:\n1. The paper claims \"Current algorithms over-fit to the training dataset and as a consequence perform poorly when deployed\nto 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 statement is still unclear to me and most of the evaluations have been on same env (i.e. learned policy tested on the env in which the dataset was collected)\n2. I am not sure how KFC++ is different from KFC", "summary_of_the_paper": "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.", "summary_of_the_review": "Weighing the strengths and concerns, I am recommending weak accept for now.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "39c2eN9zgmt", "reviewer": "Reviewer_HCMe", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I have a few questions that I believe the authors should address in their revised version to help the reader understand the method and its effectiveness.\n\n1. Can you explain whether the Koopman representation in (12) holds for systems that are not control-affine?\n2. The square is missing in the Bellman error in (1), (4), and (18). \n3. Can you explain why using the symmetry generating function in (18) is not equivalent to using a different, say non-quadratic, Bellman error? In other words, I see KFC as a method that imposes the Bellman error at states that are outside the training dataset (this is great). But since we want hat{Q}_i and Q_i to be consistent across the entire state space in order to compute the correct fixed point. Therefore, the new states in (18) can be arbitrary and do not have to come from the symmetry generating function so long as they look like legitimate transitions of the dynamical system. For example, one can compute a PCA of the local neighborhood of the samples state and use the eigenvectors to perturb the states in the dataset and constrain the Bellman error to be small there. The paper does not convince that sampling states from the Koopman operator is essential.\n4. Have you tried a standard VAE to sample new states, i.e., obtaining transitions (s_t, tilde{s_{t+1}}) where tilde{s_{t+1}} is the reconstruction of an interpolation of the latent vector corresponding to s_t and s_{t+1}? I believe this should be competitive.\n5. In the same vein, are the transformed state trajectories sensible? For instance, in (16) if one computes the new trajectory as tilde{s_t), tilde{s_{t+1}), tilde{s_{t+2}}, ..., this trajectory could be replayed in the simulator or one can compute the actions/likelihood of the actions of the learned policy for this new state trajectory to ascertain that the transformations sampled from the VAE are sensible. This is essential in order to understand the method further because the symmetry in this computation can also be trivial because (I suspect) changing the origin of the state-space is a symmetry for all environments in Table 1.\n6. This paper should study the proposed method better and conduct ablation studies. For instance, the cummulative reward of KFC should degrade as more samples are drawn from the Koopman VAE. How many extra samples are drawn to augment the existing offline datasets?\n7. Do KFC and KFC++ use the same hyper-parameters as those in CQL? The narrative in Section 4 does not seem to clarify this. If not, it is difficult to claim that the improvements in performance in Table 1 come from the new method.\n8. How many seeds were used to run the experiments in Table 1? Please provide p-values for whether KFC obtains a better cummulative reward than a baseline method.\n9. The authors of https://arxiv.org/abs/2102.09225 found that the numerical performance reported in the original CQL paper is very different from the results obtained by implementing the code of the original authors. You should cite this paper and mention whether the numbers your observations with CQL (and KFC) are consistent with those of these authors.", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper proposes an interesting idea based on Koopman theory that augments the limited dataset in offline RL but it is not clear from the current manuscript whether the performance of the method comes from the proposed innovations. The authors should present results of ablation studies that convince the reader that Koopman-operator is essential and a similar performance boost cannot be achieved by other augmentation methods (I have suggested a few above).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to improve offline RL by a data augmentation technique that exploits the symmetry of the dynamics using Koopman operator. The idea is interesting but the draft at its current form has several weaknesses as pointed out by the reviewers. The scores are borderlines at this point. I read the paper and find myself agree with reviewer ohJ3 in both the lack of  clarity and the gap in theory and empirical results. The math presentation still  a careful check and improvement. Eq(1)-(4) are already fairly confusing (should $Q_i$ and $\pi_i$ be replaced by $Q$ and $\pi$ in Eq(1)-(4), and $\hat Q$ by $\hat Q_i$ in Eq (2)-(3)?). I would like to suggest the authors to add a self-contained algorithm box for the practical algorithm procedure. Do the readers really need to understand the full Koopman theory (section 3.1) before understanding the algorithm? The authors could think about if it is better to present the practical algorithm first with minimum math, and then analyze the property of the algorithm using the math tools (and in this case, make it clear what theoretical guarantees we get exactly). I think making the paper more accessible can help the paper gain more popularity in ML readers.
Decision: Reject