K OOPMAN Q- LEARNING : O FFLINE R EINFORCEMENT L EARNING VIA S YMMETRIES OF D YNAMICS
Anonymous authors
Paper under double-blind review
A BSTRACT
Offline reinforcement learning leverages large datasets to train policies without interactions with the environment.
The learned policies may then be deployed in real-world settings where interactions are costly or dangerous.
Current algorithms over-fit to the training dataset and as a consequence perform poorly when deployed to out-of-distribution generalizations of the environment.
We aim to address these limitations by learning a Koopman latent representation which allows us to infer symmetries of the system’s underlying dynamic.
The latter is then utilized to extend the otherwise static offline dataset during training; this constitutes a novel data augmentation framework which reflects the system’s dynamic and is thus to be interpreted as an exploration of the environments phase space.
To obtain the symmetries we employ Koopman theory in which nonlinear dynamics are represented in terms of a linear operator acting on the space of measurement functions of the system and thus symmetries of the dynamics may be inferred directly.
We provide novel theoretical results on the existence and nature of symmetries relevant for control systems such as reinforcement learning settings.
Moreover, we empirically evaluate our method on several benchmark offline reinforcement learning tasks and datasets including D4RL, Metaworld and Robosuite and find that by using ourframework we consistently improve the state-of-the-art for Q-learning methods.
I NTRODUCTION
The recent impressive advances in reinforcement learning (RL) range from robotics, to strategy games and recommendation systems (Kalashnikov et al., 2018; Li et al., 2010).
Reinforcement learning is canonically regarded as an active learning process - also referred to as online RL - where the agent interacts with the environment at each training run.
In contrast, offline RL algorithms learn from large, previously collected static datasets, and thus do not rely on environment interactions (Agarwal et al., 2020a; Ernst et al., 2005; Fujimoto et al., 2019).
Online data collection is performed by simulations or by means of real world interactions e.g. robotics and in either scenario interactions maybe costly and/or dangerous.
In principle offline datasets only need to be collected once which alleviates the before-mentioned short-comings of costly online interactions.
Offline datasets are typically collected using behavioral policies for the specific task ranging from, random policies, or near-optimal policies to human demonstrations.
In particular, being able to leverage the latter is a major advantage of offline RL over online approaches, and then the learned policies can be deployed or finetuned on the desired environment.
Offline RL has successfully been applied to learn agents that outperform the behavioral policy used to collect the data (Kumar et al., 2020; Wu et al., 2019; Agarwal et al., 2020b; Ernst et al., 2005).
However algorithms admit major shortcomings in regard to over-fitting and overestimating the true state-action values of the distribution.
One solution was recently propsed by Sinha et al.
(2021), where they tested several data augmentation schemes to improve the performance and generalization capabilities of the learned policies.
However, despite the recent progress, learning from offline demonstrations is a tedious endeavour as the dataset typically does not cover the full state-action space.
Moreover, offline RL algorithms per definition do not admit the possibility for further environment exploration to refine their distributions towards an optimal policy.
It was argued previously that is basically impossible for an offline RL agent to learn an optimal policy as the generalization to near data generically leads to compounding errors such as overestimation bias (Kumar et al., 2020).
In this paper, we look at offline RL through the lens of Koopman spectral theory in which nonlinear dynamics are represented in terms of a linear operator acting on the space of measurement functions of the system.
Through which the representation the symmetries of the dynamics may be inferred directly, and can then be used to guide data augmentation strategies see Figure 1.
We further provide theoretical results on the existence on nature of symmetries relevant for control systems such as reinforcement learning.
More specifically, we apply Koopman spectral theory by: first learning symmetries of the system’s underlying dynamic in a self-supervised fashion from the static dataset, and second employing the latter to extend the offline dataset at training time by out-of-distribution values.
As this reflects the system’s dynamics the additional data is to be interpreted as an exploration of the environment’s phase space.
Some prior works have explored symmetry of the state-action space in the context of Markov Decision
Processes (MDP’s)
Higgins et al., 2018; Balaraman & Andrew, 2004; van der Pol et al., 2020) since many control tasks exhibit apparent symmetries e.g. the classic cart-pole task which is symmetric across y-axis.
However, the paradigm we introduce in this work is of a different nature entirely.
The distinction is twofold: first, the symmetries are learned in a self-supervised way and are in general not apparent to the developer; second: we concern with symmetry transformation from state tuples ( s t , s t +1 )
→ (˜ s t , ˜ s t +1 ) which leave the action invariant inferred from the dynamics inherited by the behavioral policy of the underlying offline data.
In other words we seek to derive a neighbourhood around a MDP tuple in the offline dataset in which the behavioral policy is likely to choose the same action based on its dynamics in the environment.
In practice the Koopman latent space representation is learned in a self-supervised manner by training to predict the next state using a VAE model (Kingma & Welling, 2013).
To summarize, in this paper, we propose Koopman Forward (Conservative) Q-learning (KFC): a model-free Q-learning algorithm which uses the symmetries in the dynamics of the environment to guide data augmentation strategies.
We also provide thorough theoretical justifications for KFC.
Finally, we empirically test our approach on several challenging benchmark datasets from D4RL (Fu et al., 2021), MetaWorld (Yu et al., 2019) and Robosuite (Zhu et al., 2020) and find that by using
KFC we can improve the state-of-the-art on most benchmark offline reinforcement learning tasks.
P RELIMINARIES AND BACKGROUND 2. O FFLINE RL & C ONSERVATIVE Q- LEARNING
Reinforcement learning algorithms train policies to maximize the cumulative reward received by an agent who interacts with an environment.
Formally the setting is given by a Markov decision process ( S , A , ρ, r, γ ) , with state space S , action space A , and ρ ( s t +1 | s t , a t ) the transition density function from the current state and action to the next state.
Moreover, γ is the discount factor and r ( S t ) the reward function.
At any discrete time the agent chooses an action a t ∈ A according to its underlying policy π θ ( a t | s t ) based on the information of the current state s t ∈ S where the policy is parametrized by θ .
We focus on the Actor-Critic methods for continuous control tasks in the following.
In deep
RL the parameters θ are the weights in a deep neural network function approximation of the policy or Actor as well as the state-action value function Q or Critic, respectively, and are optimized by gradient decent.
The agent i.e. the Actor-Critic is trained to maximize the expected γ -discounted cumulative reward
E π
[ (cid:80) Tt =
γ t r π ( s t , a t )] , with respect to the policy network i.e. its parameters θ .
For notational simplicity we omit the explicit dependency of the latter in the remainder of this work.
Furthermore the state-action value function Q ( s t , a t ) , returns the value of performing a given action a t while being in the state s t .
The Q-function is trained by minimizing the so called Bellman error as
[Equation]
This is commonly referred to as the i th 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
=
(cid:83) k
=1 ,..., #data-points (cid:0) s t , a t , r t , s t +1 (cid:1)
k as the option for exploration of the MDP is not available.
The policy is optimized to maximize the state-action value function via the policy improvement
[Equation]
[Figure] Note that behavioural policies including sub-optimal or randomized ones may be used to generate the static dataset D .
In that case offline RL algorithms face difficulties in the learning process.
CQL algorithm: CQL is built on top of a Soft-Actor Critic algorithm (SAC) (Haarnoja et al., 2018), which employs soft-policy iteration of a stochastic policy (Haarnoja et al., 2017).
A policy entropy regularization term is added to the policy improvement step in Eq.
(2) as
[Equation]
(3) where α either is a fixed hyperparameter or may be chosen to be trainable.
CQL reduces the overestimation of state-values - in particular those out-of distribution from D .
It achieves this by regularizing the Q-function in Eq.
(1) by a term minimizing its values over out of distribution randomly sampled action.
In the following a t +1 is given by the prediction of the policy π i ( a t +1 | s t +1 ) and ˜ α is a hyperparameter balancing the regulizer term.
The policy optimisation step is given by
[Equation]
2. K OOPMAN THEORY
Historically, the Koopman theoretic perspective of dynamical systems was introduced to describe the evolution of measurements of Hamiltonian systems (Koopman, 1931; Mezi´c, 2005).
The underlying dynamic of most modern reinforcement learning tasks is of nonlinear nature, i.e. the agents actions lead to changes of it state described by a complex non-linear dynamical system.
In contrast to linear systems which are completely characterized by their spectral decomposition non-linear systems lack such a unified characterisation.
The Koopman operator theoretic framework describes nonlinear dynamics via a linear infinite-dimensional Koopman operator and thus inherits certain tools applicable to linear control systems (Mauroy et al., 2020; Kaiser et al., 2021).
In practice one aims to find a finite-dimensional representation of the Koopman operator which is equivalent to obtaining a coordinate transformations in which the nonlinear dynamics are approximately linear.
A general non-affine control system is governed by the system of non-linear ordinary differential equations (ODEs) as
[Equation]
where s, is the n-dimensional state vector, a the m-dimensional action vector with ( s, a ) ∈ S
×
A the state-action-space.
Moreover, ˙ s = ∂
t s is the time derivative, and f is some general non-linear - at least C 1 -differentiable - vector valued function.
For a discrete time system, Eq.
(5) takes the form
[Equation]
[Equation]
where s t denotes the state at time t where F is at least C 1 -differentiable vector valued function.
Definition 1 (Koopman operator) Let K ( S × A ) be the (Banach) space of all measurement functions (observables).
Then the Koopman operator K : K ( S × A ) → K ( S × A ) is defined by
[Equation]
7) where g : S ×
A → R .
Many systems can be modeled by a bilinearisation where the action enters the controlling equationslinearly as f ( s, a ) = f
0 ( s ) + (cid:80) mi f i ( s ) a i for f i , i = 0 , . . .
, m i.e. C 1 -differentiable-vector valued functions.
In that case the action of the Koopman operator takes the simple form
[Equation]
where K ( S ) is a (Banach) space of measurement functions K 0 , K i decompose the Koopman operator in a the free and forcing-term, respectively.
Details on the existence of a representation as in equation 8 are discussed in Goswami & Paley (2017).
Associated with a Koopman operator is its eigenspectrum, that is, the eigenvalues λ , and the corresponding eigenfunctions φ λ ( s, a ) , such that [ K φ λ ]( s, a ) = λφ λ ( s, a ) .
In practice one derives a finite set of observable ⃗g = ( g 1 , . .
. , g N ) in which case the approximation to the Koopman operator admits a finite-dimensional matrix representation.
The N × N matrix representing the Koopman operator may be diagonalized by a matrix U containing the eigen-vectors of K as columns.
In which case the eigen-functions are derived by ⃗φ = U⃗g and one infers from Eq.
(5) that ˙ φ i
=
λ
i φ i , for i = 1 , . . .
, N with eigenvalues λ
i =1 ,...,N .
These ODEs admit simple solutions for their time-evolution namely the exponential functions exp( λ i t ) ..
T HE K OOPMAN FORWARD FRAMEWORK
We initiate our discussion with a focus on symmetries in dynamical control systems in Subsection 3. where we additionally present some theoretical results.
We then proceed in Subsection 3.2 by presenting the Koopman forward framework for Q-learning based on CQL.
Moreover, we discuss the algorithm as well as the Koopman Forward model’s deep neural network architecture.
Overview: Our theoretical results culminate in Theorem 3.4 and 3.5, which provides a road-map on how specific symmetries of the dynamical control system are to be inferred from a VAE forward prediction model.
In particular, Theorem 3.4 guarantees that the procedure leads to symmetries of the system (at least locally) and Theorem 3.5 that actual new data points can be derived by applying symmetry transformations to existing ones.
Practically the VAE model parametrized by a neural net is trained data from one of many behavior policies and is thus learning an approximate dynamics.
The theoretical limitations are twofold, firstly the theorems only hold for dynamical systems with differentiable state-transitions; secondly, we employ a Bilinearisation Ansatz for the Koopman operator of the system.
Practically many RL environments incorporate dynamics with discontinuous “contact” events where the Bilinearisation Ansatz may not be applicable.
However, empirically we find that our approach nevertheless is successful for environments with “contact” events. The latter does not affect the performance significantly (see Appendix D.1).
3. S YMMETRIES 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 n .
We restrict oneself to an Euclidean state space, however in general one may consider Riemannian manifolds, see e.g. Giannakis (2019).
In particular, in our work we also consider Σ invariant compact subsets of Euclidean space.
Moreover, relevant for a system of ODEs are local Lie symmetries as well as symmetries of ODEs.
Definition 2 (Local Lie Group) A parametrized set of transformations σ ϵ : S → S with s (cid:55)→ ˜ s ( s, ϵ ) for ϵ ∈ ( ϵ low , ϵ high ) where ϵ low < 0 <
ϵ high is a one-parameter local Lie group if
σ 0 is the identity map i.e for ϵ = 0 such
that ˜ s ( s, 0) =
s .
σ ϵ 1 σ ϵ 2 =
σ ϵ 2 σ ϵ 1 =
σ ϵ 1 + ϵ 2 for every | ϵ 1 | , | ϵ 2 | ≪ 1 . 3.
˜ s ( s, ϵ ) admits a Taylor series expansion in ϵ ,
i
e. in a neighbourhood of s determined by ϵ = 0 as ˜ s ( s, ϵ )
=
s + ϵ ζ ( s ) + O ( ϵ 2 ) . The points (1) and (2) in definition 2 imply the existence of an inverse element σ ϵ − 1 = σ − ϵ for | ϵ | ≪ 1 .
Moreover, note that a system of ODEs may be represented by the set of its solutions.
Definition 3 (Symmetries of ODEs) A symmetry of a system of ODEs on a locally-smooth structure (such as E n ) is a locally-defined diffeomorphism that maps the set of all solutions to itself. Definition 4 (Equivariant Dynamical System) Consider the dynamical system ˙ s = f ( s ) and let Σ be a group acting on the state-space S . Then the system is called Σ -equivariant if f ( σ · s ) =
σ · f ( s ) , for s ∈ S , ∀ σ ∈ Σ .
For a discrete time dynamical system s t +1 =
F ( s t ) one defines equivariance analogously, namely if F ( σ · s t )
=
σ · F ( s t ) , for s t ∈ S , ∀ σ ∈ Σ .
Note that as a local Lie group may satisfy the group axioms for sufficiently small parameters values it may not be a group on the entire set.
And moreover not every diffeomorphism is also an isometry.
Let us start by introducing symmetries in the simple context of dynamical system without control (Sinha et al., 2020) given by ˙ s = f ( s ) where we use analog notations as in Eq.
The Koopman operator of equivariant dynamical system is reviewed in the appendix A. Let us next turn to the case relevant for RL, namely control systems.
In the remainder of this section we focus on dynamical systems given as in Eq.
(6) and Eq.
Definition 5 (Action Equivariant Dynamical System) Let ¯Σ be a group acting on the state-actionspace S
× A of a general control system as in Eq.
(6) such that it acts as the identity operation on A i.e. σ · ( s t , a t )
= ( σ | S · s t , a t ) , ∀ σ ∈ ¯Σ .
Then the system is called ¯Σ -action-equivariant if
F ( σ · ( s t , a t )) = σ | S · F ( s t , a t ) , for ( s t , a t ) ∈ S × A , ∀ σ ∈ ¯Σ . Lemma 3.1 The map ¯ ∗ : ¯Σ × K ( S × A )
→ K ( S × A ) given by ( σ ¯ ∗ g )
( s, a ) (cid:55)−→
g ( σ − 1 · s, a ) defines a group action on the Koopman space of observables K ( S × A ) .
Theorem 3.2 A Koopman operator K of a ¯Σ -action-equivariant system s t +1 =
F ( s t , a t ) satisfies
[Equation]
In particular, it is easy to see that the biliniarisation in Eq.
(8) is Σ -action-equivariant if f i ( σ | S · s ) =
σ | S · f i ( s ) , ∀ i = 0 , . . .
, m .
Let us thus turn our focus on the relevant case of a control system s t
[Equation] =
˜ F ( s t
[Equation]
a t ) which admits a Koopman operator description as
[Equation]
where {K ( a ) } a ∈A is a family of operators with analytical dependence on a ∈ A .
Note that the bilinearisation in Eq.
(8) is a special case of Eq.
Furthermore, let { U ( a ) } a ∈A be a family of invertible operators s.t.
U ( a ) : K ( S )
→ F ( S × A ) is a mapping to the (Banach) space of eigenfunctions F ( S×A ) with eigenfunctions φ ( s, a ) :
= U ( a ) g ( s ) which obeys U ( a ) K ( a ) U ( a ) − 1 =
Λ( a ) , with Λ( a ) φ ( s, a ) =
λ φ ( a ) φ ( s, a ) and where λ φ ( a ) :
A → R .
The existence of such operators on an infinite dimensional space requires the Koopman operator in Eq.
(11) to be self-adjoint or a finite-dimensional (approximate) matrix representation to be diagonalizable.
Lemma 3.3 The map ˆ ϕ : ¯Σ × K ( S ) × { U ( a ) } a ∈A → K ( S ) given by
[Equation]
defines a group action on the Koopman space of observables K ( S ) .
Where ¯ ∗ is defined analog to
Lemma 3.1 but acting on F ( S × A ) instead of K ( S × A ) by ( σ ¯ ∗ φ )( s, a ) (cid:55)−→ φ ( σ − 1 · s, a ) .
We refer to a control system in Eq.
(11) admitting a ˆ ϕ -symmetry as ˆΣ -action-equivariant.
Theorem 3.4 Let s t {++
+} =
˜
F
( s t , a t ) be a ˆΣ -action-equivariant control system with a symmetry action as in Lemma 3.3 which furthermore admits a Koopman operator representation as
[Equation]
Then
[Equation]
[Equation] Moreover, a control system obeying equations 13 and 14 is ˆΣ -action-equivariant locally if g − 1 exists for a neighborhood of s t , i.e. then σ · ˜ F ( s t , a t ) =
˜ F ( σ · ( s t , a t )) .
Let us next provide theoretical statements on how data-points may be shifted by symmetry transformations of solutions of ODEs.
To establish an easier connection to the next section we introduce the notation Let E : S → K ( S ) and D : K ( S ) → S denote the C 1 -differentiable encoder and decoder to and from the finite-dimensional Koopman space approximation, respectively, i.e. E ◦ D
=
D ◦ E = id .
Theorem 3.5 Let s t +1 =
˜ F ( s t , a t ) be a control system as in equation 13 and σ a t an operator obeying equation 14.
Then σ ϵa t : (cid:0) s t , s t +1 , a t (cid:1) (cid:55)−→ (cid:0) ˜ s t , ˜ s t +1 , a t (cid:1) with
[Equation]
(15) is a one-parameter local Lie symmetry group of ODEs. In other words one can use a symmetry transformation to shift both s t as well as s t +1 such that ˜ s t
+1 =
˜ F (˜ s t , a t ) .
3. KFC ALGORITHM
In the previous section we laid the theoretical foundation for the KFC-algorithm by providing a raod-map on how to derive symmetries of dynamical control systems based on a Koopman latent space representation.
The goal is to generate new data points for the RL algorithm at training-time as
Eq.
(15) in Theorem 3.5.
The reward r t is not part of the symmetry shift process and will just remain unchanged as an assertion.
On the power of using symmetries: Let us emphasize the practical advantage of employing symmetries to generate augmented data points.
It is evident from Theorem 3.5 that a symmetry transformation shifts both s t as well as s t +1 which evades the necessity of forecasting states.
Thus the use of an inaccurate fore-cast model is avoided and the accuracy and generalisation capabilities of the VAE are fully utilized.
The magnitude of the induced shift is controlled by the parameter ϵ ≪ 1 such that | s − ˜ s | =
O ( ϵ ) to limit out-of-distribution generalisation errors . Algorithmically the symmetry maps are derived in two distinct ways which we denote KFC and KFC++ . The latter, constitute a simple starting point to extract symmetries from our setup. Limitations: Theorems 3.4 and 3.5 imply that if an operator commutes with the Koopman operator we can find a local symmetry group.
However, global conditions may not be easily inferred. Moreover, in practice one has that D ◦ E ≈ id . Thus all assumptions of theorem 3.5 hold approximately.
The Koopman forward model: The KFC algorithm requires pre-training of a Koopman forward model F : S × A → S which is closely related to a VAE architecture as [Equation] where both of E and D are approximated by Multi-Layer-Perceptrons (MLP’s) and the bilinear
Koopman-space operator approximation are implemented by a single fully connected layer for K i , i = 0 , . . . , m , respectively.
The model is trained on batches of the offline data-set tuples ( s t +1 , s t , a t ) and optimized via an additive loss-function of the VAE and the forward prediction part of the model.
We refer the reader to Appendix C for details.
We integrate our framework into a specific Q-learning algorithm (CQL). In practice the Koopman latent space representation is N-dimensional (finite). The Koopman operator and the symmetry generators admit matrix representations; matrix multiplication replaces the mapping ˆ ∗ .
Following in the footsteps of (Sinha et al., 2021)
our approach leaves the policy improvement Eq.
(3) unchanged but modifies the policy optimisation step Eq.
(4) as
[Equation]
The state-space symmetry generating function σ ϵa : S → S depends on the normally distributed random variables ϵ .
Note that we only modify the Bellman error of Eq.
(17) and leave the CQL specific regulizer untouched.
We study two distinct cases which differ on an algorithmic level
[Equation]
where σ a t ( ⃗ϵ ) = R e (cid:0) U ( a t ) diag ( ϵ 1 , . . .
, ϵ N ) U − 1 ( a t ) (cid:1) , (18) and U ( a t ) diagonalizes the Koopman operator with eigenvalues
λ
i , i = 1 , . . .
, N i.e. the latter can be expressed as K ( a t ) =
U ( a t ) diag ( λ 1 , . . . , λ N ) U − 1 ( a t ) .
Note that we abuse our notation as E ( s ) ≡ ⃗g ( s ) =
[ g 1 ( s ) , . . .
, g N ( 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 σ a t to commute with the approximate Koopman operator in Eq.
For the KFC option in Eq.
(18) the symmetry generator matrix σ a is obtained by solving the equation σ a t · K ( a t ) − K ( a t ) · σ a t = 0 which may be accomplished by employing a Sylvester algorithm (syl).
For KFC++ we compute the eigen-vectors of K ( a ) which then constitute the columns of U ( a ) .
Thus in particular one infers that [ σ a t ( ⃗ϵ ) , K ( a t )]
= 0 .
The latter, is solved by construction of σ a ( ⃗ϵ ) which commutes with the Koopman operator for all values of the random variables.
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 p K ; otherwise we use a random normally distributed state shift.
E MPIRICAL 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), we perform experiments on 6 different challenging robotics tasks from
MetaWorld (Yu et al., 2019) and RoboSuite (Zhu et al., 2020).
We compare KFC to the baseline
CQL algorithm (Kumar et al., 2020), and two best performing augmentation variants from S4RL,
S4RL- N and S4RL-adv (Sinha et al., 2021).
We use the exact same hyperparameters as proposed in the respective papers.
Furthermore, similar to S4RL, we build KFC on top of CQL (Kumar et al., 2020) to ensure conservative Q-estimates during for policy evaluation.
4. 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 on various different types of environments suggests that KFC and KFC++ fundamentally improves the data augmentation strategies discussed in S4RL.
KFC-variants also improve the performance of learned agents on challenging environments such as antmaze: which requires hierarchical planning, kitchen and adroit tasks: which are sparse reward and have large action spaces.
Similarly, KFC-variants also perform well on difficult data distributions such as “medium-replay”: which is a collected by simply using all the data that the policy encountered while training base SAC policy, and “-human” which is collected using human demonstrations on robotic tasks which results in a non-Markovian behaviour policy (more details can be found in the D4RL manuscript (Fu et al., 2021)).
Furthermore, to our knowledge, the results for KFC++ are state-of-the-art in policy learning from D4RL datasets for most environments and data distributions.
We do note that S4RL outperforms KFC on the “-random” split of the data distributions, which is expected as KFC depends on learning a simple dynamics model of the data to use to guide the data augmentation strategy.
Since the “-random” split consists of random actions in the environment, our simple model is unable to learn a useful dynamics model. For ablation studies see appendix D.
[Table]
medium-expert 111.
walker-random 7.
walker-medium 79.
walker-medium-replay 21.
walker-medium-expert 98.
Adroit pen-human 37.
pen-cloned 39.
hammer-human 4.
hammer-cloned 2.
door-human 9.
door-cloned 0.
relocate-human 0.
relocate-cloned -0.
Franka kitchen-complete 43.
kitchen-partial 49.
[Figure]
4. M ETAWORLD AND R OBOSUITE BENCHMARKS
To further test the ability of KFC, we perform additional experiments on challenging robotic tasks.
Following (Sinha et al., 2021), we perform additional experiments with 4 MetaWorld environments (Yu et al., 2019) and 2 RoboSuite environments (Zhu et al., 2020).
We followed the same method to collect the data as described in Appendix F of S4RL (Sinha et al., 2021), and report the mean percent of goals reached, where the condition of reaching the goal is defined by the environment.
We report the results in Figure 2, where we see that using KFC to guide the data augmentation strategy for a base CQL agent, we are able to learn an agent that performs significantly better.
Furthermore, we see that for more challenging tasks such as “push” and “door-close” in the MetaWorld, KFC++ outperforms the base CQL algorithm and the S4RL agent by a significant margin.
These set of experiments further highlight the ability of KFC to guide the data augmentation strategy.
R ELATED 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 in Sinha et al.
However, the latter randomly augments the data while our augmentation framework is based on symmetry state shifts.
Regarding the connection to world models (Ha & Schmidhuber, 2018).
Here a VAE is used to decode the state information while a recurrent separate neural network predicts future states.
Their latent representation is not of Koopman type.
Also no symmetries and data-augmentations are derived.
Algebraic symmetries of the state-action space in Markov Decision Processes (MDP) originate (Balaraman & Andrew, 2004) an were discussed recently in the context of RL in (van der Pol et al., 2020).
Their goal is to preserve the essential algebraic homomorphism symmetry structure of the original MDP while finding a more compact representation.
The symmetry maps considered in our work are more general and are utilized in a different way.
Symmetry-based representation learning (Higgins et al., 2018) refers to the study of symmetries of the environment manifested in the latent representation.
The symmetries in our case are derived form the Koopman operator not the latent representation directly.
In (Caselles-Dupré et al., 2019) the authors discuss representation learning of symmetries (Higgins et al., 2018) allowing for interactions with the environment.
A Forward-VAE model which is similar to our Koopman-Forward VAE model is employed.
However, our approach is based on theoretical results providing a road-map to derive explicit symmetries of the dynamical systems as well as their utilisation for state-shifts.
In (Sinha et al., 2020) the authors extend the Koopman operator from a local to a global description using symmetries of the dynamics.
They do not discuss action-equivariant dynamical control systems nor data augmentation.
In (Salova et al., 2019) the imprint of known symmetries on the blockdiagonal Koopman space representation for non-control dynamical systems is discussed.
This is close to the spirit of disentanglement (Higgins et al., 2018).
Our results are on control setups and deriving symmetries.
On another front, the application of Koopman theory in control or reinforcement learning has also been discussed recently.
For example, Li et al.
(2020) propose to use compositional
Koopman operators using graph neural networks to learn dynamics that can quickly adapt to new environments of unknown physical parameters and produce control signals to achieve a specified goal.
Kaiser et al.
(2021) discuss the use of Koopman eigenfunction as a transformation of the state into a globally linear space where the classical control techniques is applicable.
To the best of our knowledge, this paper is the first to discuss Koopman latent space for data augmentation.
C ONCLUSIONS
In this work we proposed a symmetry-based data augmentation technique derived from a Koopman latent space representation.
It enables a meaningful extension of offline RL datasets describing dynamical systems, i.e. further "exploration" without additional environment interactions.
The approach is based on our theoretical results on symmetries of dynamical control systems and symmetry shifts of data.
Both hold for systems with differentiable state transitions and with a Bilinearisation
Ansatz for the Koopman operator.
However, the empirical results show that the framework is successfully applicable beyond those limitations.
We empirically evaluated our method on several benchmark offline reinforcement learning tasks D4RL, Metaworld and Robosuite and find that by using our framework we consistently improve the state-of-the-art of Q-learning algorithms.
A K OOPMAN THEORY AND S YMMETRIES
This section provides some additional introductory material on Koopman theory and symmetries of dynamical systems.
Lemma A.1 The map ∗ : Σ × K ( S )
→ K ( S ) given by ( σ ∗ g )( s ) (cid:55)−→ g ( σ − 1 · s ) defines a group action on the Koopman space of observables K ( S ) .
Theorem A.2 Let K be the Koopman operator associated with a Σ -equivariant system s t +1 =
F ( s t ) .
Then
[Equation]
Theorem A.2 states that for a Σ -equivariant system any symmetry transformation commutes with the Koopman operator.
For the proof see (Sinha et al., 2020).
To proceed let us re-evaluate some information from the main text of this work.
Let { U ( a ) } a ∈A be a family of invertible operators s.t.
U ( a ) : K ( S )
→ F ( S × A ) is a mapping to the (Banach) space of eigenfunctions φ ( s, a ) :
=
U ( a ) g ( s ) ∈ F ( S × A ) .
Which moreover obeys U ( a ) K ( a ) U ( a ) − 1 = Λ( a ) , with Λ( a ) φ ( s, a ) =
λ φ ( a ) φ ( s, a ) and where λ φ ( a ) :
A → R .
The existence of such operators puts further restriction on the Koopman operator in Eq.
However, as our algorithm employs the finite-dimensional approximation of the Koopman operator i.e. its matrix representation this amounts simple for K ( a ) to be diagonalizable and U ( a ) is the matrix containing its eigen-vectors as columns.
To evolve a better understanding on the required criteria for the infinite-dimensional case we employ an alternative formulation of the so called spectral theorem below.
Theorem A.3 (Spectral Theorem) Let K be a bounded self-adjoint operator on a Hilbert space H .
Then there is a measure space ( S , Σ , µ ) and a real-valued essentially bounded measurable function λ on S and a unitary operator U : H
→ L 2 µ ( S ) i.e. U ∗ U
=
U U ∗ =
id such that
[Equation]
[Equation]
In other words every bounded self-adjoint operator is unitarily equivalent to a multiplication operator.
In contrast to the finite-dimensional case we need to slightly alter our criteria to U ( a ) K ( a ) U ( a ) − 1 = Λ( a ) , with Λ( a ) φ ( s, a ) =
λ φ ( a, s ) φ ( s, a ) and where λ φ ( s, a ) : S ×
A → R .
Concludingly, a sufficient condition for our criteria to hold in terms of operators on Hilbert spaces is that the Koopman operator K ( a ) is self-adjoint i.e. that
[Equation]
B P ROOFS
In this section we provide the proofs of the theoretical results of Section 3.
{+Proof of Lemma A.1 and Theorem A.
+}
The proofs of the lemma as well as the theorem can be found in (Sinha et al., 2020).
Proof of Lemma 3.1: We aim to show that map ¯ ∗ :
¯Σ × K ( S × A ) → K ( S × A ) given by ( σ ¯ ∗ g )( s, a ) (cid:55)−→ g ( σ − 1 · s, a ) defines a group action on the Koopman space of observables K ( S ×A ) , where ¯Σ defines the symmetry group of definition 5.
Firstly, let g ∈ K
( S × A ) and σ 0 ∈ ¯Σ be the identity element, then we see that it provides existence of an identity element of ¯ ∗ by
[Equation]
Secondly, let σ 1 , σ 2 ∈ ¯Σ and ⊙ denoting the group operation i.e. σ 1 ⊙ σ 2 = σ 3 ∈ ¯Σ .
Then (cid:16)
[Equation]
where in ( I ) we have used the invertibility of the group property of ¯Σ .
Lastly it follows analogously that for σ, σ − 1 ∈ ¯Σ that
[Equation]
Thus the existence of an inverse is established which concludes to show the group property of ¯ ∗ .
Proof of Theorem 3.2: We aim to show that with K being the Koopman operator associated with a ¯Σ -action-equivariant system s t +1 = F ( s t , a t ) . Then [Equation]
First of all note that by the definition of the Koopman operator of non-affine control systems it obeys
[Equation] Using the latter one thus infers that [Equation] where in (I) we have used that it is a ¯Σ -action-equivariant system.
Moreover, one finds that
[Equation]
which concludes the proof.
Proof of Lemma 3.3: We aim to show that the map ˆ ϕ : ¯Σ × K ( S ) × { U ( a ) } a ∈A → K ( S ) given by
[Equation]
s [Equation] , defines a group action on the Koopman space of observables K ( S ) .
Where ¯ ∗ is defined analog to
Lemma 3.1 but acting on F ( S × A ) instead of K ( S × A ) by ( σ ¯ ∗ φ )( s, a ) (cid:55)−→ φ ( σ − 1 · s, a ) .
First of all note that
[Equation]
We proceed analogously as in the proof of Lemma 3.1 above.
Firstly, let g ∈ K ( S ) and σ 0 ∈ ¯Σ be the identity element then on infers from Eq.
(22) that it provides existence of an identity element of ˆ ϕ by
[Equation]
where we have used the notation
[Equation]
Secondly
[Equation]
let σ 1 , σ 2 ∈ ¯Σ and ⊙ denoting the group operation
i
[Equation]
e. σ 1 ⊙ σ 2 = σ 3 ∈ ¯Σ .
Then
[Equation]
where in ( I ) we have used the invertibility of the group property of Σ .
Lastly, it follows analogously that for σ, σ − 1 ∈ ¯Σ that
[Equation]
Thus the existence of an inverse is established which concludes to show the group property of ˆ ϕ .
Proof of Theorem 3.4: We aim to show twofold.
= ⇒ :
Firstly, that with s t +1 =
˜
F
( s t , a t ) be a ˆΣ -action-equivariant control system with a symmetry action as in Lemma 3.3 which furthermore admits a Koopman operator representation as
[Equation]
Then
[Equation]
⇐ = :
Secondly, the converse.
Namely, if a control system s t +1 =
˜ F ( s t , a t ) obeys Eqs.
(13) and (14), then it is ˆΣ -action-equivariant, i.e. σ · ˜ F ( s t , a t ) =
˜ F ( σ · ( s t , a t )) .
For notational simplicity we drop the subscripts referring to time in the remainder of this proof i.e. s t → s and a t → a .
Let us start with the first implication i.e. = ⇒ .
First of all note that by the definition of the Koopman operator one has
[Equation]
Using the latter one infers that
[Equation]
Moreover, one derives that
[Equation]
which concludes the proof of the first part of the theorem.
Let us next show the converse implication i.e. ⇐ = .
For this case it is practical to use the discrete time-system notation explicitly.
Let σ ∈ Σ be a symmetry of the state-space and be ˜ s t = σ · s t and ˜ s t +1 =
σ · s t +1 the σ -shifted states.
Then
[Equation]
Thus in particular
[Equation]
where in (I) we have used that the symmetry operator commutes with the Koopman operator.
Moreover, one finds that
[Equation]
from which one concludes that
[Equation]
Finally, we use the invertibility of the Koopman space observables i.e. g − 1 ( g ( s ))
=
s to infer
[Equation]
However, in general s the Koopman space observables are not invertible globally. The "inverse function theorem" guarantees the existence of a local inverse if g ( s ) is C 1 differentiable for maps between manifolds of equal dimensions. However, we assume inevitability locally.
[Equation]
which at last concludes our proof of the second part of the theorem.
Extension of Theorem 3.5: Moreover, one may account for practical limitations i.e. an error by the assumption [ σ ϵa t , K ( a t )]
= ϵ a 1 . One then finds that ˜ s t +1 = ˜
F
(˜ s t , a t )
+ O ( ϵ a ) .
Thus the error becomes suppressed when ϵ a ≪ ϵ .
The error may be due to practical limitations of capturing the true dynamics as well as the symmetry map.
Proof of Theorem 3.5: We will show here theorem 3.5 as well as the extension mentioned in the paragraph prior to this proof.
{+Let s t +
+}
= ˜ F ( s t , a t ) be a ˆΣ -action-equivariant control system as in
Theorem 3.4.
For symmetry maps
[Equation]
we aim to show that one can use a symmetry transformation to shift both s t as well as s t +
[Equation]
[Equation]
From equation 24 and the definition of the Koopman operator one infers that
[Equation]
By using that [ σ a t , K ( a t )] = 0 and equation 24 one finds that [Equation]
[Equation]
which concludes the proof of the first part of the theorem with σ a t → ( 1 + ϵ σ a t ) .
Local diffeomorphism.
Per definition the maps D, E are differentiable and invertible which implies that they provide a local diffeomorphism from and to Koopman-space.
Also the linear map i.e. a matrix multiplication by ( 1 + ϵ σ a t ) is a diffeomorphism.
Thus the symmetry map σ ϵa t i.e. Eq.
(15) constitutes a local diffeomorphism.
The above proof culminating in Eq.
(24) implies that we have a local diffeomorphism mapping solutions of the ODEs to solutions, thus a symmetry of the system of
ODEs.
Limitations.
However D, E to be invertible is a strong assumption as it implies that the Koopman space approximation admit the same dimension as the state space.
Note that however in practice as we only require D ◦ E ≈ i
d one may choose other Koopman space dimensions.
Local Lie group.
What is left to show is that the symmetry of the system of ODEs locally is a Lie group.
In definition 2 we need to show points (1)-(3).
For ϵ =
one finds that σ 0 a t is the identity map i.e for such that
[Equation]
[Equation]
[Equation]
for every | ϵ 1 | , | ϵ 2 | ≪ 1 .
˜ s t ( s t , ϵ ) admits a Taylor series expansion in ϵ , i.e. in a neighbourhood of s determined by ϵ = 0 .
We may Taylor expand D around the point E ( s t ) ≡ g as
[Equation]
30) thus ˜ s
[Equation]
( s
[Equation]
31) for ζ ( s t ) =
(cid:80) N I =1 ( σ a t E ( s t ))
I ∂D ∂g
I .
The existence of an inverse element ( σ ϵa t ) − 1 = σ − ϵ a t 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 [ σ ϵa t , K ( a t )] = ϵ a 1 one finds ˜ s t +1 =
˜
F
(˜ s t , a t ) + O ( ϵ a ) . Note that σ a t → ( 1 + ϵ σ a t ) admits an expansion in the parameter ϵ . For ϵ ≪one can Taylor expand the differentiable functions D, E to find that
[Equation]
where g σ = σ a t ˆ ∗ E ( s t ) and g = E ( s t ) , and with the index I = 1 , . . .
, N .
The analog expression holds for t + 1 i.e. ˜ s t {++
 =+} s t +1 + O ( ϵ ) .
If a t = a t ( s t ) is differentiable function of the states 7 one may Taylor expand the Koopman operator using equation 32 as
[Equation]
where and a ′ in ( s t ) =
∂a t ∂ st,n | s t and n = 1 , . .
, dim ( s t ) and with
[Equation]
However, by the definition of the symmetry map a t (˜ s t ) = a t ( s t ) , thus the Koopman operators in equation 33 match ∀ ϵ .
By using the assumption on the error one infers from equation 26 that
[Equation]
One may Taylor expand D by making use of ϵ a ≪ 1 as
[Equation]
Thus under the assumption that the numerical error is given by a violation of the commutation relation one finds
[Equation]
By comparison with equation 32 one infers that the ϵ -expansion is a good approximation to the real dynamic if ϵ a ≪ ϵ .
This concludes the proof of theorem 3.5.
Let us note that the proof for state shifts arising from a sum of symmetry transformations as
[Equation]
is completely analogous, and an analogous theorem holds.
We choose the number of symmetries to run over the dimension of the Koopman latent space N as it is relevant for the KFC++ setup.
C I MPLEMENTATION DETAILS
The Koopman forward model The KFC algorithm requires pre-training of a Koopman forward model F : S ×
A → S as
[Equation]
if c = 0 corresponds to a VAE.
[Equation]
c
[Equation]
[Equation]
prediction model [Equation] [Equation] where both of E and D are approximated by Multi-Layer-Perceptrons (MLP’s) and the bilinear
Koopman-space operator approximation are implemented by a single fully connected layer for
K i =0 ,...,m , respectively.
The model is trained on batches of the offline data-set tuples ( s t +1 , s t , a t ) and optimized via the loss-function
[Equation]
where s ϵ is a normally distributed random sate vector shift with zero mean and a standard deviation of 6 · 10 − 2 where the ¯ l 2 -loss is the Huberloss.
For training the Koopman forward model in Eq.
(16) we split the dataset in a randomly selected training and validation set with ratios 70% / 30% .
We then loop over the dataset to generate the required symmetry transformations and extend the replay buffer dataset with the latter according to Eq.
(18) as
[Equation]
It is more economic to compute the symmetry transformation as a pre-processing step rather than at runtime of the RL algorithm.
Let us reemphasize a couple of points addressed already in the main text to make this section selfcontained.
Note that we have built upon the implementation of CQL (Kumar et al., 2020), which is based on SAC (Haarnoja et al., 2018).
We also use the same hyperparameters for all the experiments as presented in the CQL and S4RL papers for baselines.
Koopman forward VAE-model:
For the encoder as well as the decoder we choose a three layer
MLP architecture, with hidden-dimensions [512 , 512 , N ] and [512 , 512 , m ] , respectively.
Where
N = 32 is the Koopman latent space dimension and m is the state space dimensions following our notation from section 2.
The activation functions are ReLU (Nair & Hinton, 2010).
We use an ADAM optimizer with learning rate 3 · 10 − 4 .
We train the model for 75 epochs.
The batch size is chosen to be 256 .
Lastly, we employ a random normally distributed state-augmentation for the VAE training i.e. s → s + s ϵ with s ϵ ∈ N (0 , 1 · 10 − 2 ) .
Hyper-parameters KFC:
We take over the hyper-parameter settings from the CQL paper (Kumar et al., 2020) except for the fact that we do not use automatic entropy tuning of the policy optimisation step Eq.
(3) but instead a fixed value of α = 0 .
The remaining hyper-parameters of the conservative
Q-learning algorithm are as follows: γ = 0 .
99 and τ = 5 · 10 − 3 for the discount factor and target smoothing coefficient, respectively.
Moreover, the policy learning rate is 1 · 10 − 4 and the value function learning rate is 3 · 10 − 4 for the ADAM optimizers.
We enable Lagrange training for ˜ α
with threshold of 10 . 0 and learning rate of 3 · 10 − 4 where the optimisation is done by an ADAM optimizer.
The minimum Q-weight is set to 10 .
and the number of random samples of the Q-minimizer is 10 .
For more specifics on these quantities see (Kumar et al., 2020).
The model architectures of the policy as well as the value function are three layer MLP’s with hidden dimension 256 and ReLU -activation.
The batch size is chosen to be 256 .
Moreover, the algorithm performs behavioral cloning for the first 40 k training-steps, i.e. time-steps in an online RL notation.
The KFC specific choices are as follows: the split of random to Koopman-symmetry state shifts is 20 / 80 i.e. p K = 80% , whereas the option for the symmetry-type is either "Eigenspace" or "Sylvester" referring to case (I) and (II) of Eq.
(18), respectively.
Regarding the random variables in the symmetry transformation generation process, we choose ϵ i ∈ N (0 , 1 · 10 − 4 ) for case (II) and ϵ ∈ N (0 , 5 · 10 − 5 ) for case (I) in Eq.
(18) after normalizing σ a by its mean.
While for the random shift we use ˜ ϵ ∈ N (0 , 3 · 10 − 3 ) which is the hyper-parameter used in (Sinha et al., 2021).
Computational setup: We performed the empirical experiments on a system with PyTorch 1.9.0a (Paszke et al., 2019)
The hardware was as follows: NVIDIA DGX-
with 16 V100 GPUs andcores of Intel(R)
Xeon(R)
Platinum 8168 CPUs and NVIDIA DGX-
with 8 A100 GPUs withcores of Intel(R)
Xeon(R) E5-2698 v4 CPUs.
The models are trained on a single V100 or A100 GPU [Table]
D A BLATION STUDY
[Table]
D. T RAINING ON NON CONTACT EVENTS
To test the theoretical limitations in a practical setting one interesting experiment is to use the S4RLN augmentation scheme when the current state is a “contact” state and to choose KFC++ otherwise.
We define a “contact” state as one where the agent makes contact with another surface; in practice we looked at the state dimensions which characterize the velocity of the agent parts, and if there is a sign difference between s t and s t +1 , then contact with a surface must have been made.
The details regarding which dimension maps to agent limb velocities is available in the official implementation of the Open AI Gym suite (Brockman et al., 2016).
Note that the theoretical limitations are twofold, firstly the theorems 3.4 and 3.5 only hold for dynamical systems with differentiable state-transitions; secondly, we employ a Bilinearisation Ansatz for the Koopman operator.
Contact events lead to non-continuous state transitions and moreover would require incorporating for external forces in our Koopman description, thus both theoretical assumptions are practically violated in our main empirical evaluations of KFC and KFC++.
By performing this ablation study which is more suitable for our Ansatz, i.e. the KFC++ part is only trained on "differentiable" trajectories without contact events.
In other words, the Koopman description is unlikely to give good information regarding s t +
if there is a “contact event”, we simply use S4RL- N when training the policies, which we expect to give slightly better performance than the KFC++ baseline.
We report the results in Table 2, where we perform results on the OpenAI Gym subset of the D4RL benchmark test suite.
We denote the proposed scheme as KFC++-contact and compare to both
S4RL- N and KFC++.
We see that KFC++-contact performs comparably and slightly better in some instances, compared to KFC++.
This is expected as contact events may be rare, or the data augmentation given by the KFC++ model during contact events is pseudo-random and itself is normally distributed, which may suggest that KFC++ behaves similarly to KFC++-contact.
However, we do not say this conclusively.
D. M AGNITUDE OF KFC AUGMENTATION
To further try to understand the effect of using KFC for guiding data augmentation, we measure the mean L2-norm of the strength of the augmentation during training.
We measured the mean L2-norm between the original state and the augmented state, and found the distance to be 2 .
6 × 10 − 4 .
For comparison, S4RL uses a fixed magnitude of 1 × 10 − 4 (Sinha et al., 2021) for all experiments with
S4RL-Adv.
Since S4RL-Adv does a gradient ascent step towards the direction with the highest change, its reasonable for the magnitude of the step to be smaller in comparison.
An adversarial step too large may be detrimental to training as the step can be too strong.
KFC minimizes the need for making such a hyperparameter search decision. D. S YMMETRIES
In this section we provide a quantitative analysis of the symmetry shifts generated by KFC++ in comparison to random shifts of the latent states.
For the latter we choose the variance such that
absolute state shift
{+[Equation] averaged over all sampled data points (D4RL) is comparable between the two setups. Moreover we simply use the encoder We then compare how well the shifted states are in alignment with the actual dynamic using the online Mujoco environment . To do so we set the environment state to the value ˜ s t and use the action a t to perform the environment step which we denote by s t +
+}
=
M ( s t , a t )
The performance metric is then given by
[Equation]
Note that ∆ E is simply the error to the actual dynamic of the environment.
We take the model trained on the hopper-medium-expert dataset from section appendix D.1 to compute the symmetry shift. Moreover, for the comparison in equation 43 we split the state space into positions and velocities of the state vector, i.e. the first five elements of s t +1 − M (˜ s t , a t ) give position while the remaining ones given velocities. The norm is taken subsequently. For the results see figure D.3.
We conclude that there is a qualitative difference between the distributions obtained. Most notable the symmetry shift allows for a controlled shift by a larger ∆ S while still maintaining a small error ∆ E . This amounts to a wide but accurate exploration of the environments state-space. Although we cannot say this conclusively we expect this to be the reason for the performance gains of KFC++ over random shifts.
D. P REDICTION M ODEL
As a further ablation study it is of interest to compare the symmetry induced state shifts tuple (˜ s t , ˜ s t +1 ) to one obtained by the forward prediction of our VAE-forward model equation 16.
[Figure]
KFC++prediction: In particular we use the forward prediction on ˜ s t obtained by the KFC++ shift, i.e. ˜ s t +1 = F c =1 (˜ s t , a t ) .
We refer to this setup in the following as KFC++prediction.
See tablefor the results.
Fwd-prediction: Moreover, we study the case where we ˜ s t is obtained by a shift with a normally distributed random variable N (0 , 6 · 10 − 3 ) and then simply forward predict as ˜ s t +1 = F c =1 (˜ s t , a t ) .
This is comparable to conventional model based approaches employing our simple VAE forward model. This constitutes an ideal systematic comparison as the model, hyper-parameters and training procedure are identical to the one used in the KFC variants. Moreover, the variance of the random variable is such that the distance to the augmented states is comparable to the ones in KFC, see appendix D.2. See table 4 for the results.
We conclude that KFC++prediction falls behind both KFC as well as KFC++ .
This was expected as the symmetry shift in the latter alters the original data points by means of the VAE which admits not only a much higher accuracy but also advanced generalisation capabilities to out of distribution values.
Moreinteresting however is that Fwd-prediction falls behind both KFC as well as KFC++ as well as KFC++prediction .
This is a strong indicator that the symmetry shifts provide superior out-of-distribution data for training a Q-learning algorithm.
D. D ISCUSSION
Although our KFC framework is suited best for the description of environments described by control dynamical systems the learned (in self-supervised manner) linear Koopman latent space representation may be applicable to a much wider set of tasks.
However, there are notable shortcomings to the current implementation both conceptually as well as practically.
The bilinearisation in Eq.
(16) of the latent space theoretically assumes the dynamical system to be governed by Eq.
(8), which is rather restrictive.
Although a general non-affine control system admits a bilinearisation (Brunton et al., 2021) [Table] [Table] it generically requires the observables to depend on the action-space variables implicitly.
Secondly, the Koopman operator formalism is theoretically defined by its action on an infinite dimensional observable space.
The finite-dimensional approximation i.e. the latent space representation of the
Koopman forward model in Eq.
(16) lacks accuracy due to that.
On the practical side our formalism requires data pre-processing which is computationally expensive, i.e. solving the Sylvester or eigenvalue problem for every data-point.
Moreover, the Koopman forward model in Eq.
(16) serves as a function approximation to two distinct tasks.
Thus one faces a twofold accuracy vs. over-estimation problem which needs to be balanced.
The systematic error in the VAE directly imprints itself on the state-data shift in Eqs.
(15) and (18) and may thus conceal any potential benefit of the symmetry considerations.
Lastly, the dynamical symmetries do not infer the reward.
Thus the underlying working assumption is that the reward should not vary much in Eq.
A note on the simplicity of the current algorithm: Let us stress a crucial point.
Algorithmically the symmetry maps are derived in two distinct ways KFC and KFC++ .
The latter, constitute a simple starting point to extract symmetries from our setup.
More elaborate studies employing the extended literate on Koopman spectral analysis are desirable.
Moreover, it is desirable to extend our framework to more complex latent space descriptions such as e.g. world models.
Both on a theoretical as well as a practical level.
It is our opinion that by doing so there is significant room for improvement both in the accuracy of the derived symmetry transformations as well their induced performance gains of Q-learning algorithms.
Note that currently our VAE model is of very simple nature and the symmetries are extracted in a rather uneducated way.
While the Sylvester algorithm simply converges to one out of many symmetry transformations for the KFC++ algorithm we omit all the information of the imaginary part, let alone utilize concrete spectral information.