Efficient Dynamics Modeling in Interactive Environments with Koopman Theory
Authors: Arnab Kumar Mondal
Siba Smarak Panigrahi
Sai Rajeswar
Kaleem Siddiqi
Siamak Ravanbakhsh
Venue: ICLR 2024 poster (2024)
Keywords: Koopman Theory
Reinforcement Learning
Dynamical System
Planning
Longe range dynamics prediction models
Efficient forward dynamics
Primary area: reinforcement learning
TL;DR: We propose to use a Diagonalized Koopman operator to model Interactive environments efficiently in the latent space as linear dynamical systems with a control input.

Abstract: The accurate modeling of dynamics in interactive environments is critical for successful long-range prediction. Such a capability could advance Reinforcement Learning (RL) and Planning algorithms, but achieving it is challenging. Inaccuracies in model estimates can compound, resulting in increased errors over long horizons.
We approach this problem from the lens of Koopman theory, where the nonlinear dynamics of the environment can be linearized in a high-dimensional latent space. This allows us to efficiently parallelize the sequential problem of long-range prediction using convolution while accounting for the agent's action at every time step.
Our approach also enables stability analysis and better control over gradients through time. Taken together, these advantages result in significant improvement over the existing approaches, both in the efficiency and the accuracy of modeling dynamics over extended horizons. We also show that this model can be easily incorporated into dynamics modeling for model-based planning and model-free RL and report promising experimental results.

planning and Published as a conference paper at ICLR 2024
EFFICIENT DYNAMICS MODELING IN INTERACTIVE EN- VIRONMENTS WITH KOOPMAN THEORY
Arnab Kumar Mondal Mila, McGill University
Siba Smarak Panigrahi Mila, McGill University
Sai Rajeswar ServiceNow Research
Kaleem Siddiqi Mila, McGill University
Siamak Ravanbakhsh Mila, McGill University
ABSTRACT
The accurate modeling of dynamics in interactive environments is critical for suc- cessful long-range prediction. Such capability could advance Reinforcement Learning (RL) and Planning algorithms, but achieving it is challenging. Inaccu- racies in model estimates can compound, resulting in increased errors over long horizons. We approach this problem from the lens of Koopman theory, where the nonlinear dynamics of the environment can be linearized in a high-dimensional latent space. This allows us to efficiently parallelize the sequential problem of long-range prediction using convolution while accounting for the agent's action at every time step. Our approach also enables stability analysis and better control over gradients through time, Taken together, these advantages result in significant improvement over the existing approaches, both in the efficiency and the accuracy of modeling dynamics over extended horizons. We also show that this model can be easily incorporated into dynamics modeling for model-free RL and report promising experimental results.
INTRODUCTION
The ability to predict the outcome of agent's action over long horizons is crucial unresolved
challenge in Reinforcement Learning (RL) (Sutton & Barto 2018: Mnih et al. (005)] Silver et al.
2017|Mnil et al., 2016).. This is especially important in model-based RL and planning, where deriving a policy from the learned dynamics models allows one to efficiently accomplish a wide
variety of tasks in environment (Du & Narasimhan) 2019 Hafner et al. 2020 Sikehi al. 2021
Hansen et al.] 2022; Schrittwieser et al. 20200 et al.] 2022). In fact, state-of-the-an model-free techniques also rely on dynamics models to learn a better representation for downstream value prediction tasks (Schwarzer et al., 2020). Thus, obtaining accurate long-range dynamics models in the presence of input and control is crucial. In this work, we leverage techniques and perspectives from Koopman theory (Koopman| 1931 Mauroy et al.) 2020 Koopman & Neumann 1932 Brunton et al.) 2021> to address this key problem in long-range dynamics modeling of interactive environments. The application of Koopman theory allows us to linearise nonlinear dynamical system by creating a bijective mapping to linear dynamics a possibly infinite dimensional space of observables I Conveniently, deep neural network can learn to produce such a mapping, enabling reliable approximation of the non- non-linear dynamics in the finite dimension a linear latent space so that only finite subset of the most relevant coamplex-valuede Koopman observables need to be tracked.
We show that this linearization has two major benefits: 1) operator are directly related to the stability and expressive power of the controlled dynamics model. Furthermore, spectral Koopman theory provides a unifying perspective over several existing directions on stabilizing gradients through time, including long-range sequence modeling using state-space models (Gu et al.) 2020 2022b Gupta et al. (002).. and the use of unitary matrices (Arjovsky al. 2016) latent dynamic models; 2) This can help avoid a computational bottleneck due
The eigenvalues of the linear latent-space
"Correspondence: monabmonndleee @ miilaquuebec. Work done while at ServiceNow Research
'The term observables can be misleading as it refers the latent space in the machine learning jargon.

---

Published as a conference paper at ICLR 2024
Non-linear latent
MLP-based dynamics model
dynamics
88
Complex dyramical system
d5
Model Long Horizons
Parallal and fast prediction
Control gradients while BPTT
8,41
Koopman- based dynamics model
Simple linear latent dynamics!
8 Kix
Figure 1: A comparison of our Koopman-based linear dynamics model with non-linear ML.P-basd dynamics model. The Diagonal Koopman formulation allows for modeling longer horizons efficiently with control over gradients. Here BPTT stands for Backpropagation Through Time.
to the sequential nature of dynamics. This is achieved by a reformulation using convolution and diagonalization of the linear operator. In other words, one perform efficient parallel training of the model over time steps, despite dealing with controlled dynamical system.
Our experimental results in offlineRRL datasets demonstrate the effectiveness of approach for reward and state prediction over long horizon. In particular, we report competitive results against dynamics modeling baselines that use Multi-Layer Perceptrons (MLPs), Gated Recurrent Units (GRUs), Transformers and Diagonal State Space Models Gu et al. (2022a) while being significantly faster. Finally, we also present encouraging results for model-bases planning and model-free RL with our Koopman-basee dynamics model.
In summary, our contributions are:
1. We formulate an efficient linear latent dynamics model from Koopman theory that can be parallelized during training for long-range dynamics modeling in interactive environments.
2. We demonstrate the use of the spectral decomposition of the Koopman matrix, which enables better control over gradients through time and accelerated training.
3. We show that our technique serves as preferred alternative for existing dynamics models in planning and RL algorithms, offering enhanced performance, speed and training stability.
PRELIMINARIES
2. KOOPMAN THEORY FOR DYNAMICAL SYSTEMS
In the context of non-linear dynamical systems, the Koopman operator, a.k.a. the Koopman-von Neumann operator, linear operator for studying the system's dynamics. The Koopman operator is defined as linear transformation that acts on the space of functions or observables of the system, known the observables space F. For non-linear discrete or continuous time dynamical system
dæ f(x)
Fe+ F(I,) or dl
the Koopman operator F + F, is defined as Kgg g o F where is the set of all functions observables that form infinitediimensional Hilbert space. In other words, for every function g X + R belonging to F. where I, E X C R", we have (aC) o (F(())) g(x441)-
2

---

Published as a conference paper at ICLR 2024
The infinite dimensionality of the Koopman operator presents practical limitations. Addressing this, we seek to identify a subspace g C F, spanned by a finite set of observables : where typically m > n, that approximates invariance under the Koopman operator.
Constraining the Koopman operator on this invariant subspace results in a finiteddimensional linear operator K E Cm called the Koopman matrix, which satisfies g = K (g())) Usually, the base observation functions {9:} i are hand-crafted using knowledge of the system's 's underlying physics. However, data-driven methods have recently been proposed learn the Koopman operator by representing the base observations or their eigenfunction basis using deep neural networks, where decoder reconstructs the input from linear latent dynamics (e.g.. [uucch al.) [018 Champion et al.) 2019). Lincarizing the dynamics allows for closed- solutions for the predictions of the dynamical system, as well as stability analysis based on the eigendecomposition of the Koopman matrix Fan et al. (2022) Yi & Manchester (2023).
The observation functions ) spanning an invariant subspace can always be chosen to eigenfunc- tions { of the Koopman operator, i.e., Ko (a)) = (x)) Vx. With this choice of observation
functions, the Koopman matrix becomes diagonal,
KD diag(>i....... An), A vi.
By using diagonal Koopman matrix, are effectively offloading the task of learning the suitable eigenfunctions of the Koopman operator that form an invariant subspace to the neural network encoder. Additionally, for continuous time input, be able model higher-order frequencies in the eigenspectrum and provide better approximations, one could adaptively choose the most relevant eigenvalues and eigenfunctions for the current state (Lusch al. (018). This means g(æ++1)= = KD (D(M(((z)))g(zz) where À : G + C" can be a neural network.
2.2 APPROXIMATE KOOPMAN WITH CONTROL INPUT
The Koopman operator can be extended to non-linear control systems, where the state of the sys- tem is influenced by external control input u, such that Iz-1 F(x,, u,) or de de f(x, u). Simply treating the pair (z, u) as the input F, the Koopman operator becomes (Kg) u,) = g ((F(2 (2 ue), u4-1) = g( u4+1). If the effect of the control input on the system's 's dynamics is linear, i.e. the control affine setting, we have
f(x, u) fo(x) +
M(x)u
assuming finite number of eigenfunctions. For such control-affine system, could show that the Koopman operator is bilinearized (Brunton et al. 2021 Bruder et al. 2021):
g(#+41) K(u)g(xx) (Ko +
us ())((()))
More generally, could make the Koopman operator function of ue that the resulting Koopman matrix satisfies g(x+41) = K ((x)) This is the approach taken by Weissenbacher al. (2022) to model symmetries of dynamics in offline RL. When combined with the diagonal Koopman matrix of Eq. ( and fixed modulus [Ag] = 1'vi, the Koopman matrix becomes a product of ()) representations, and the resulting latent dynamics model, resembles a special setting of the symmetry-based approach of Mondal et al. (2022). where group representations are used to encode states and dete-dependent actions.
An alternate approach to approximating the Koopman operator with control signal assumes a decou- pling of state and control observables g(x, u) [=(x), ff [u)] [g (ar),... Im (z), fi (u), f..,,...f... (u))
Brunton et al. (2021): Bruder et al. (2019). This gives rise to a simple linear evolution:
g( = K g((())) xx)) + (f
(2)
where K E Cmxm and L E mm are matrices representing the linear dynamics of the state observables and control input, respectively. We build this approach in our proposed methodology. Although, as shown in Brunton et al. (2016) 202T); Bruder et al. (2019). even assuming a linear control term Lag perform well in practice we use neural networks for both f and g. Our
3

---

Published as a conference paper at ICLR 2024
rationale for choosing this formulation is that the additive form of Eq. combined with the diagonalized Koopman matrix of Eq. (T), enable fast parallel training across time-steps using a convolution operation where the convolution kernel can be computed efficiently. Moreover, this setting is amenable to the analysis of gradient behaviour, as discussed in Section [32.1] For a more comprehensive overview of Koopman theory in the context of controlled and uncontrolled dynamics, we direct readers to Brunton et al. (2021); Mauroy et al. (2020).
PROPOSED MODEL
Roopnan Dynanic ca Model
Boopnan bynanica ca Block
action embeddiing
$
State Ehonder
& C-ne Diagonal Koopman
Lec
Dynamis
Next slate embeoding
emtecding
Figure A schematic of the latent Koopman dynamics model. Both actions and initial state embedding are encoded into latent space in complex (C) domain before passing through the Koopman dynamics block. (see Appendix for an efficient Jax implementation of the model)
3.1 LINEAR LATENT DYNAMICS MODEL.
Our task in dynamics modeling is to predict the sequence of future states St+1, .... St+ given a starting state 8; and some action sequence og, 0++1, ..., at+--1, assuming a Markovian system. Following the second approach described in Section 2.2] we assume that state and control observables are decoupled, and use neural network encoder to encode both states and actions
Tz (8)) and 41 fa(
Due to this decoupling, the continuous counterpart of latent space dynamics Eq. (2) is dr dt
= KI((t) + Lu(t),
(4)
and the solution is given by z(s) = =x x0)+ (((-t) Lu Lu(t)dt where ex is a matrix exponential. One could discretize this to get ât+1 = Kx + Lu,, where K and L are obtained by Zero-Order Hold (ZOH) discretization the continuous-time equation (Iserles 2009):
K = (tK)
L = ((KK- (K) Atk) 1)L
In practice, assume that observations sampled uniformly in time and use learnable time step parameter At. Unrolling the discrete version of the dynamics for 7 steps, we get:
I
K
K
0
I
K
[1+1. 1........8 "r.,,"",,,44+1 00,77
[K.... ""T,..",,jaaa KK]a + [a,, -
(5)
where C++. Luc+ee encodes the effect of each input on all future observables, and the iin distinguishes the prediction from the ground truth observable æ. (see Appendix B
We can then recover the predicted state sequence by inverting the function ge, using decoder network 8r4k de (f++2). Loss functions in both input and latent space can be used to learn the encoder,
We CNN for pixel input and ML.P for state-based input and actions.
4

---

Published as a conference paper at ICLR 2024
decoder, and latent dynamics model. We refer to these loss functions the Koopman consistency loss and the state cante-preiction loss, respectively:
Lonnsistnccy
Batt Ie+all2
Latate pred
Ide( 8
(6)
3.2 DIAGONALIZATION EFFICIENCY AND STABILITY OF THE KOOPMAN OPERATOR
Since the set of diagonalizable matrices is dense in C and has full measure, we can always diagonalize the Koopman matrix as shown Eq. (T). Next, we show that this diagonalization can be leveraged for efficient and parallel forward evolution and training.
To unroll the latent space using Eq. (5) we need perform matrix exponentiation and dense multi- plication. Using a diagonal Koopman matrix K = diag(Ai..... .. Am), this calculation is reduced to computing X (r + 1) complex- complex-valued valued Vandermonde matrix A, where At.j XT and doing row-wise circular convolution of this matrix with a sequence of vectors.
Recall that E C complex vectors. We use superscript index their elements i.e., It [e...... x²"]. With this notation, the i-th element the predicted Koopman observables for future time steps is given by (see Appendix for the derivation):
Fi41 (4+r) [ (.. + [1,
X
where @ is circular convolution with zero padding. The convolution can be efficiently computed for longer time steps using Fast Fourier Transform. (Brigham 1988) (see Appendix
3.2. GRADIENTS THROUGHI TIME
We show how we can control the behavior of gradients through time by constraining the real part of the eigenvalues of the Diagonal Koopman matrix. Let μ and refer the real and imaginary part of the Koopman eigenvalues - that is A = P.j + iwj.
Theorem 3.1. For every time step
{I,...T} the discrete dynamics, the norm of the gradient of any loss k-step given by Ls with respect latent representation at time step given by xa is scaled version of the norm of the gradient of the same loss by 24+k. where the scaling factor depends on the exponential of the real part of the Koopman eigenvalues, that is:
0L* Oz7
0Lk dz
= 2d0ps)
Vj E (1,.,,),
and similarly, for all S k the norm of the gradient of Ls with respect to the control input time step 1+1-1 given by is 6++1-g a scaled version of the norm of the gradient Ls by T14k, where the scaling factor depends on the exponential of the real part of the Koopman eigenvalues, that is:
0Lk e1(()-
0Lk
Eje (1,.,))
dc 0ct
D& 1
proof of this theorem is provided in Appendix A The theorem implies that the amplitude of the gradients from a future time step scales exponentially with the real part of each diagonal Koopman eigenvalue. We use this theorem for better initialization the diagonal Koopman matrix as explained in the next section.
3.2.2 INITIALIZATION OF THE EIGENSPECTRUM
The imaginary part of the cigenvalues of the diagonal Koopman matrix captures different frequency modes of the dynamics. Therefore, is helpful to initialize them using increasing order frequency, constant a, to cover wide frequency range.
that is, W ajæ, for
we know that the real part of the eigenvalues impacts the gradient's behavior through time. To avoid vanishing gradients and account for prediction errors over longer horizons, one could eliminate the real part H := This choice turns the latent transformations into blocks
From Theorem A.1
of 2D rotations (Mondal et al. 2022 This is also related to using Unitary Matrices to avoid

---

Published as a conference paper at ICLR 2024
State Prediction ErroP
Tanaformer
Tranfhorme
GRU
+
Koepman
Koooman
Kospman
Reward Prediction Eroor
NUP
Fransformer
Tariformer
Trarsfomerr
GRU
CRJ
CRU
Kocpman
DSSM
Kooumaa
DS5H
6.0
Kappiaan
DESM
616
expert
mediumeespee
fullreglay
exgert
medumeexpet
fubreppay
espert
meplummeeeeet
fulleggey
(a) hoppor-v2
(b) walker2d-v2
(c) halfchcotah-v2
Figure 3: Forward state and reward prediction error in Offline Reinforcement Learning environments. We consider five dynamics modeling techniques and perform this prediction task over a horizon of 100 environment steps. The results over runs. Our Koopman-based method is competitive with the best performing GRU baseline while being 2x faster. See Appendix for exact numerical values.
vanishing exploding gradients in Recurrent Neural Networks (Arjovsky et al. 2016). However, intuitively, we might prefer to prioritize closer time steps. This can be done using small negative values e.g., M4 E (-0.,, 0.2, 0.33. An alternative having fixed real part to turn into
a bounded learnable parameter μy E 0.3, -0.1]. We empirically found (4j == -0.2 Vj, and
wj == j be good choices and use this as our default initialization. In Appendix we report the
results of an ablation study with different initialization schemes.
EXPERIMENTS
4.1 LONG-RANGE FORWARD DYNAMICS MODELING WITH CONTROL
We model the non-linear controlled dynamics of
Training Speed Comparison
MuJoCo environments (Todorov et al.) 2012)
Training Speed (tterationssecome
GRU
using our Koopman operator. We choose standard MLP-based latent non-linear dynam- ics model with two linear layers followed by ReLU -ineariy as one of our baselines. To make it easier compare with Koopman dynam- ies model, we use state embeddings of the same dimensionality using Go. which MIP, as de- scribed in Section 3.1 This embedding is then
240
Transformer
120
200-
80-
00-
40-
23-
fed to the MLP-based latent dynamics model or the diagonal Koopman operator. This approach is widely used in RL for dynamics modeling Hansen et al.) 2022 Schrit- twieser et al. 2020 Schwarzer et al. 2020). To
Training Sequence Length
Figure 4: Training speed iterations/second (t) for the state prediction task using different dy- namics model on hal Fcheetah--exerrt--.. Each iteration consists of one gradient update of the entire model using a mini-batch of 256 A00 GPU. See Table for exact numerical values.
(Sikchi al. 2021
further compare our model with more expres- sive alternatives that can be trained in parallel, design causal Transformer (Vaswani al. 2017. Chen et al. 2021) based dynamics model, which takes a masked sequence of state-actions and outputs representations that are used to predict next states and rewards Following the success of Hafner et al. 22200), we also design a GRU-based dynamics model that also takes a masked sequence
Alll the after the starting are masked, but actions fed the model for all future time-steps.

---

Published as a conference paper at ICLR 2024
of state-actions output representations for state and reward prediction. Finally, we the same strategy to also design a DSSM-based Gu et al. (2022a) dynamics model. To ensure a fair comparison in terms of accuracy and runtime, maintain equal number of trainable parameters for all the aforementioned models. (see Appendix [ for more details)
For the forward dynamics modeling experiments, a popular offline-RL RL environment. We select offline datasets trajectories collected from three different popular Gym-MuloCo control tasks: hopper, walker, and halfcheetale These trajectories are obtained from three distinct quality levels of policies, offering range of data representing various levels of expertise in the task: expert, medhum-expert, and full replay. We divide the dataset of 1M samples into 80:20 splits for training and testing, respectively. To train the dynamics model, we randomly sample trajectories of length T from the training data, where T the horizon specified during training. We test our learned dynamics model for horizon length of 100 by randomly sampling 50,000 trajectories of length 100 from the test set. We use our learned dynamics model predict the ground truth state sequence given the starting state and sequence of actions.
the D4RL Fu al. (2020) dataset, which
Training stability We can only train the MLP-based dynamics model for 10-time steps across all enviromments. Using longer sequences often results in exploding gradients and instability in training due to backpropagation through time (Sutskever 2013). The alternative of using tanh in MLP-base models can lead to vanishing gradients. contrast, our diagonal Koopman operator handles long trajectories without encountering vanishing or exploding gradients, in agreement with our discussion in Section 3.2.1
4.1.1 STATE AND REWARD PREDICTION
In Fig. we evaluate our model's 's accuracy in state and reward prediction over long horizons involving 100 environment steps. For this experiment, we add reward predictor and jointly minimize the reward prediction loss and the state prediction loss. We set the weight of the consistency loss in the latent space to 0.001. We see that for longer horizon prediction, our model considerably better in predicting rewards and states accurately for hopper, walker, and hal f fcheetah in comparison to MLP, Transformer and DSSM based model while being competitive with GRU-based model. Furthermore, Fig. 4 empirically verifies that proposed Koopman dynamics model is significantly faster than an MI.,, GRU, DSSM and Transformer based dynamics model. The results also reveal the trend that model's relative speed-up over baselines grows with the length of the horizon.
Mean
IQM
Optimality Gap
MLP
Kaapman
0.792 0.70 Human Normalized Score
0.30
Figure 5: Comparison of Koopman-based dynamics model (with a horizon of 20) and MLP- based dynamics model of vanilla TD-MPC (Hansen et al., 2022). The results are over 5 random seeds for each environment. Higher Mean & IQM and lower Optimality Gap is better.
Mean
IQM
Optimality Gap
Vanilla
Koopman
0.660 0.675
0.35
Human Normalized Score
Figure 6: Comparison of vanilla SAC Haarnoja al. 2018) and its integration with an based (SPR) and a Koopman-based dynamics model for incorporating self-predictive representations in the DeepMind Control Suite. The results are over 5 random seeds for each environment. Higher Mean & IQM and lower Optimality Gap is better.
7

---

Published as a conference paper at ICLR 2024
4.2 KOOPMAN DYNAMICS MODEL FOR RL AND PLANNING
We now present promising results from integrating the diagonal Koopman model into two areas: (I) model-based planning, and (II) model-free free reinforcement learning (RL). In the latter, the dynamics model enhances representation learning. Both approaches aim to solve continuous control tasks in data-efficient manner. We provide a brief summary of these methods and their modifications with the Koopman dynamics model in Appendix [ While the same dynamics model can also be used for (III) model-based RL, we leave that direction for future work.
To assess the performance of our algorithm, we calculate the average episodic return at the conclusion of training and normalize it with the maximum score, i.c., 1000. However, due to significant variance over different runs, relying solely on mean values may not provide reliable metrics. To address this, Agaawal et al. (202T suggests using bootstrapped confidence intervals (CI) with stratified sampling. which is particularly suitable for small sample sizes, such as the five runs per environment in our case. By utilizing bootstrapped CI, we obtain interval estimates indicating the range within which an algorithm's aggregate performance is believed to fall and report the Interquartile Mean (IQM). The IQM represents the mean across the middle 50% of the runs, providing robust measure of central tendency. Furthermore, we calculate the Optimality Gap (OG). which quantifies the extent to which the algorithm falls short of achieving normalized score of 1.0.
Evaluation Metric
4.2.1 MODEL-BASED PLANNING
We utilize TD-MPC (Hansen et al. 2022) as the baseline for our moodel-based planning approach. TD-MPC combines the benefit of model-free approach in long horizons by learning a value function while using an MI.--aased "Task-Oriented Latent Dynamics" (TOLD) model for planning over shorter horizons. To assess the effectiveness of our proposed dynamics model, we replace TOLD with our diagonal Koopman operator. Subsequently, we trained the TD-MPC agent from scratch using this modified Koopman TD-MPC approach. See Appendix C.4 for details on this adaptation. To evaluate the performance of our variation, we conducted experiments in four distinct environ- ments: Quadruped Run, Quadruped Walk, Cheetah Run, and Acrobot Swingup. We run experiments in the state space and compare them against vanilla TD-MPC.
Our Koopman dynamics model was trained with a horizon of 20-time steps. We observe that TD-MPC suffers from training instability, performance degradation and high variance when using 20-time steps. Consequently, we opted to use horizon of 5-time steps for the vanilla approach. 5 suggests that the Koopman TD-MPC outperforms the MLP-based baseline while being more stable.
4.2.2 MODEL-FREE RL
Current data-efficient model-free RL approaches like (SPR: (PR:Waewarrerr et al., 2020 use dynamics modeling as auxiliary task to improve representation learning. We apply our dynamics model to this model-free RI. framework, where we use an adaptation of SPR as our baseline. This adaptation uses an MI.P-based dynamics model rather than the original CNN used in SPR in order to eliminate the advantage of the original SPR for image inputs. SPR uses an auxiliary consistency loss similar to Eq. (6), as well as augmentations and other techniques from Grill et al. (2020) to prevent repre- sentation collapse. It also improves the encoder's ' representation quality by making it invariant to the transformations that are not relevant to the dynamics using augmentation; see also (Srinivas et al.) : Yarats et al.] (202T). To avoid overconstraining the latent space using the dynamics model, SPR uses projection heads. Similarly, we encode the Koopman observables using a projection layer above the representations used for Q-learning We provide more details on the integration of the Koopman dynamics model into the SPR framework in Appendix C.5
We perform experiments on five different environments from the DeepMind Control (DMC) Suite (Tunyasuvunakool al. 2020 Tassa al.) 2018) (Ball in Cup Catch, Cartpole Swingup, Cheetah Run, Finger Spin, and Walker Walk) in pixel space, and com- pare with an MLP-based dynamics model (SPR) and no dynamics model, i.e., SAC (Haarnoja et al.) 2018) as baselines. Figure demonstrates that for horizon of 5 time steps doopman-based model outperforms the SAC plus M..P-based dynamies model baseline while being more efficient.
'A smaller Optimality Gap indicates superior performance.
8

---

Published as a conference paper at ICLR 2024
5 RELATED WORK
Diagonal state space models (DSSMs). DSSMs (Gu et al., 2022al Gupta et al. 2022 Mchta
et al. 2022 have been shown to be an efficient alternative to Transformers for long-range sequence modeling. DSSMs with certain initializations have been shown [uu et al. (022aa) to approximate convolution kernels derived for long-range memory using the HiPPO (Gu al. 2020) theory. This explains why DSSMs can perform as well as structured state space models (S4) (Gu et al., (022bb) as pointed out by Gupta al. (2022). Inspired by the success of gating in transformers (Hua et al.. 2022), MMhtt et al. (2022) introduced a gating mechanism in DSS to increase its efficiency further. However, a DSSM a non-linear model used for sequence modeling and functions differently from our Koopman- bassed dynamics model, as explained in Appendix D
Probabilistic latent spaces. Several early papers use approximate inference with linear and non- linear latent dynamics so as to provide probabilistic state representations (Haarnoja et al., 2016) Becker et al., 2019 Shaj et al.) 2021, Hua et al., 2022 Fraccaro al., (017). A common theme in several of these articles is the use of Kalman filtering. using variational inference or MCMC, for estimating the posterior over the latents and mechanisms for disentangled latent representations. While these methods vary in their complexity, scalability, and representation power, similar to an LSTM or a GRU, they employ recurrence and, therefore, cannot be parallelized.
Koopman Theory for Control. The past work in this area has solely focused on learning the Koopman operator or discovering representations for Koopman spectrum for control (Shi & Meng 2022, Watter et al., 2015). Korda & Mezić (2018) extends the Koopman operator to controlled dynamical systems by computing a finite-dimensional approximation of the operator. Furthermore, they demonstrate the of these predictors for model predictive control (MPC). |an et al. (2020)
tasks data-driven approach to use neural networks to represent the Koopman operator in controlled dynamical system setting. Kaiser al. (20TT) introduces a method using Koopman eigenfunctions for effective linear reformulation and control of nonlinear dynamical systems. [Song et al. (2021)
proposes deep Koopman reinforcement learning that learns control tasks with a small amount of data. However, none of these works focus on the training speed and stability for predicting longer horizon.
6 CONCLUSION
Dynamics modeling has a key role to play in sequential decision-makee at training time, having access to a model can lead sample-efficient training, and at deployment, the model can be used for planning. Koopman theory is an attractive foundation for building such a dynamics model in Reinforcement Learning. The present article shows that models derived from this theoretical foundation can be made computationally efficient. We also demonstrate how this view gives insight into the stability of learning through gradient descent in these models. Empirically, we show the remarkable effectiveness of such an approach in long-range dynamics modeling and provide some preliminary yet promising results for model- based planning and model-free RL.
LIMITATIONS AND FUTURE WORK
Although our proposed Koopman- based dynamics model has produced promising results, our current treatment has certain limitations that motivate further investigation. Our current model is tailored for deterministic environments, overlooking stochastic dynamics modeling. We intend to expand our research in this direction, based on developments around stochastic Koopman theory (Mezić 2005) Wanner & Mezic 2022). where Koopman observables become random variables, enabling uncertainty estimation in dynamics. Secondly, although our state prediction task yielded impressive results, we have identified areas for further growth, particularly in applications to RL. and planning.
We hypothesize that the distribution shift during training and the objective change in training and evaluation hurt the performance of our Koopman dynamics model. We also plan to conduct a more comprehensive study involving a wider range of tasks and environments. Additionally, we aim to explore the compatibility of our model with different reinforcement learning algorithms to showcase its adaptability. Finally, believe that the diagonal Koopman operator holds promise for model- based RL, where robust dynamics modeling over long horizons is precisely what is currently missing. We intend to investigate this direction, potentially unlocking avenues for advancement.

---

Published as a conference paper at ICLR 2024
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep reinforcement learning the edge of the statistical precipice. Advances neural information processing systems, 11-23304-29320 2021.
Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International conference machine learning, pp. 1220-128. PMLR, 2016.
Philipp Becker, Harit Pandya, Gregor Gebhardt, Cheng Zhao, C James Taylor, and Gerhard Neumann. Recurrent kalman networks: Factorized inference in high-dimensional deep feature spaces. In International conference on machine learning, pp. 544-552. PMLR, 2019.
Oran Brigham. The fast Fourier transform and its applications. Prentice-Hall, Inc., 1988.
Daniel Bruder, Brent Gillespie, C David Remy, and Ram Vasudevan. Modeling and control of soft robots using the koopman operator and model predictive control. arXiv preprint arXiv:1902.0280 002.09277..227, 2019.
Daniel Bruder, Xun Fu, and Ram Vasudevan. Advantages of bilinear koopman realizations for the modeling and control of systems with unknown dynamics. IEEE Robotics and Automation Letters, 6(3):4369-431)) 2021.
Steven L Brunton, Joshua L Proctor, and J Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 113(15):3932-399) -9322, 2016.
Steven I. Brunton, Marko Budišić, Eurika Kaiser, and Nathan Kutz. Modern koopman theory for dynamical systems. arXiv preprint arXiv:2102. 12086, 2021.
Kathleen Champion, Bethany Lusch, Nathan Kutz, and Steven I. Brunton. Data-driven discovery of coordinates and governing equations. Proceedings of the National Academy of Sciences, 116(45): 22445-22451, 2019.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 14:15004-15097, 2021.
Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In International conference on computers and games, pp. 72-83. Springer, 2006.
Yilun Du and Karthic Narasimhan. Task-agnostic dynamics priors for deep reinforcement learning. In International Conference on Machine Learning, pp. 1696-1705. PMLR, 2019.
Fletcher Fan, Bowen Yi, David Rye, Guodong Shi, and Ian Manchester. Learning stable koopman embeddings. 2022 American Control Conference (ACC), pp. 2742-2747. IEEE, 2022.
Marco Fraccaro, Simon Kamronn, Ulrich Paquet, and Ole Winther. A disentangled recognition and nonlinear dynamics model for unsupervised learning. Advances in neural information processing systems, 30, 2017.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 33:21271-21288. 2020.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474-1487, 2020.
10

---

Published as a conference paper at ICLR 2024
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initial- ization of diagonal state space models. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022a. URL ht.tps: http::://openrevieee ew. net /ortt//ormmm d=yJE7 iQSeep
Albert Gu. Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022b. URL https
//openreview net./ forum? d dwu v1AC
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. Advances in Neural Information Processing Systems, 11122982-229992 2022.
Tuomas Haarnoja, Anurag Ajay, Sergey Levine, and Pieter Abbeel. Backprop kf: Learning discrim- inative deterministic state estimators. Advances in neural information processing systems, 29,
2016.
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 Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp- 1861-1870. PMLR, 10-15 Jul 2018. URL https://proceediing mir press httrees///80/haaanoi....... htm
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference Learning Representations, 2020.
URL htttps://openreeeee wwwp:e/wwwweeeeeooroetteeeeeee
net font:/ffrmm??dde 1tt
Yiqiang Han, Wenjian Hao, and Umesh Vaidya. Deep learning of koopman representation for control. In 2020 59th IEEE Conference on Decision and Control (CDC), pp. 1890-1895. IEEE Press, 2020.
Nicklas A Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predictive control. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 8387-8406. PMLR, 17-23 Jul 2022. URL https: https::/prroceeding. mlx wwweess/w122haasee222222 html
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International Conference Machine Learning, pp. 9099-9117. PMLR, 2022.
Arieh Iserles. A first course in the numerical analysis of differential equations. Number 44. Cambridge university press, 2009.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations, 2017. URL
Lett/Eorun:id-SJyyyD55
Arnay Kumar Jain, Shivakanth Sujit, Shruti Joshi, Vincent Michalski, Danijar Hafner, and Samira Ebrahimi Kahou. Learning robust dynamics through variational sparse gating. Advances in Neural Information Processing Systems, 35:1612-1626, 2022.
Xiaoyu Jiang, Qiuxuan Chen, Shiyi Han, Mingxuan Li, Jingyan Dong, and Ruochen Zhang. When to trust your model: Model-based policy optimization, 2020. httthhtts:///ppeee net/ (etoorummmiissgPPiPee
ew Submitted NeurIPS 2019 Reproducibility Challenge.
Eurika Kaiser, J Nathan Kutz, and Steven Brunton. Data-driven discovery of koopman eigenfunc- tions for control. Machine Learning: Science and Technology, 2(3):035023, 2021.
Bernard O Koopman. Hamiltonian systems and transformation in hilbert space. Proceedings of the National Academy of Sciences, 1755):35--318, 1931.
Bernard O Koopman and v Neumann. Dynamical systems of continuous spectra. Proceedings of the National Academy of Sciences, 18(3):255-263. 1932.

---

Published as a conference paper at ICLR 2024
Milan Korda and Igor Mezić. Linear predictors for nonlinear dynamical systems: Кooрman operator meets model predictive control. Automatica, 93; 11:1449660, 2018. ISSN 0005-1098, doi: https: sci
URL https:/ http::/wwwwwww.eeeeieic.ee 8301 33x]
com/
Timothy P. Lillicrap. Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Yoshua Bengio and Yann LeCun (eds.), ICLR, 2016.
Kendall Lowrey, Aravind Rajeswaran, Sham Kakade, Emanuel Todorov, and Igor Mordatch. Plan online, learn offline: Efficient learning and exploration via model-based control. In International Conference on Learning Representations, 2019. WtLhttts::///oo
id=Byey'n05FQ
Bethany Lusch, Nathan Kutz, and Steven L Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nature communications, 9(1):1-10, 2018.
Alexandre Mauroy, Yoshihiko Susuki, and Igor Mezić. Introduction to the koopman operator in dynamical systems and control theory. The koopman operator in systems and control: concepts, methodologies, and applications, pp. 3-33, 2020.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces, arXiv preprint arXiv:2206.139477 2022.
Igor Mezić. Spectral properties of dynamical systems, model reduction and decompositions. Nonlin- ear Dynamics, 41:309-325, 2005.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller. Playing atari with deep reinforcement learning. CoRR, 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattic, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):559-5533 2015.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap. Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. ICMI., 2016.
Arnab Kumar Mondal, Vineet Jain, Kaleem Siddiqi. and Siamak Ravanbakhsh. Eqr: Equivariant representations for data-efficient reinforcement learning. In International Conference on Machine Learning, pp. 15908-15926. PMLR, 2022.
Masashi Okada and Tadahiro Taniguchi. Variational inference mpc for bayesian model-based reinforcement learning. In Leslie Pack Kaelbling, Danica Kragic, and Komei Sugiura (eds.), Proceedings of the Conference on Robot Learning, volume 100 of Proceedings of Machine Learning Research, pp. 258 272.22722222 PMLR, 30 Oct-01 Nov 2020. URL htlthttts///pcceeeeet mlr. 1aaa.pppes///100//aaaaa....
Reuven Rubinstein and Dirk P Kroese. The cross-entropy method: a unified approach to com- binatorial optimization, Carte-Carro-aa simulation, and machine learning, volume 133. Springer, 2004.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604 609, 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv, abs/1707 0aw/1777.06333 000700637, 2017.
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bach- man. Data effaccffic reinforcement learning with self preficteeee representations. In International Conference on Learning Representations, 2020.
12

---

Published as a conference paper at ICLR 2024
Vaisakh Shaj, Philipp Becker, Dieter Büchler, Harit Pandya, Niels van Duijkeren, James Taylor, Mare Hanheide, and Gerhard Neumann. Action- conditional recurrent kalman networks for forward and inverse dynamics learning. In Conference Robot Learning, pp. 765-781, PMLR, 2021.
Haojie Shi and Max Q.-H. Meng. Deep koopman operator with control for nonlinear systems. IEEE Robotics and Automation Letters, 2022.
Harshit Sikehi, Wenxuan Zhou, and David Held. Learning off-policy with online planning. In 5th Anmual Conference on Robot Learning, 2021. URL
[d=1GNV9SW95ed
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioamnis Antonoglou, Veda Panneershelvam, and Mare Lanctot. Mastering the game of go with deep neural networks and tree search. Nature, 550:354 359, 2017.
Lixing Song, Junheng Wang, and Junhong Xu, A data-efficient reinforcement learning method based on local koopman operators. In 2021 20th IEEE International Conference on Machine Learning and Applications (ICMIA), 2021.
A. Srinivas, Michael Laskin, and P. Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, 2020.
Ilya Sutskever. Training recurrent neural networks. University of Toronto Toronto, ON, Canada, 2013.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint 1rXiv:I801.000 1801.00690, 2018
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 JEFE/RSJ International Conference on Intelligent Robots and Systems, 2012.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm_control: Software and tasks for continuous control. Software Impacts, 6:100022, 2020. ISSN 2665-9638. doi: https://doi. 2020. 2020. 100022, URL https: /www. scwencedirect......cc.c......... wtw.eaiitettctceceoeeeooeaeeeeee /oonsscience/ atticlee/si//222533332333
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Mathias Wanner and Igor Mezic. Robust approximation of the stochastic koopman operator. SIAM Journal on Applied Dynamical Systems, 1931),1 2022.
Manuel Watter, Jost Springenherg, Joschka Boedecker, and Martin Riedmiller. Embed to control: locally linear latent dynamics model for control from raw images. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances in Neural Information Processing Systems, 2015.
Matthias Weissenbacher, Samarth Sinha, Animesh Garg, and Kawahara Yoshinobu. Koopman q- learning: Offline reinforcement learning via symmetries of dynamics. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 23645-23667 PMLR, 17-23 Jul 2022. URL http:htttss:ppeeeeeeeeeeeeeeeeee html
m1 mlr. pr.press///.... / v 62 wei htt.eeesss..eeeeeeeeeeeaaaaaaaaa
Grady Williams, Andrew Aldrich, and Evangelos Theodorou. Model predictive path integral control using covariance variable importance sampling, 2015. URL https: //tp:////xxxi./// org.bbs/.rg/////.//// 1509. 01149
13

---

Published as a conference paper at ICLR 2024
Grady Williams, Paul Drews, Brian Goldfain, James M Rehg, and Evangelos A Theodorou. Information-theoree model predictive control: Theory and applications to autonomous driv- ing. IEEE Transactions Robotics, 34(6):1603-1622 2018.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations, 2021. URL URLattps: https:/
Bowen Yi and Ian R Manchester. On the equivalence of contraction and koopman approaches for nonlinear stability and control. IEEE Transactions on Automatic Control, 2023,
14

---

Published as a conference paper at ICLR 2024
A PROOF OF THEOREM 3.1
Theorem A.1. For every time step k E {),...} in the discrete dynamics, the norm of the gradient of any loss at k-step given by Lk with respect to latent representation at time step t given by is scaled version of the norm of the gradient of the same loss by 2 where the scaling factor depends on the exponential of the real part of the Koopman eigenvalues, that is:
0L* = 0x
ŌLk
Vj E ml....m.
DrL+
and similarly, for all l S the norm of the gradient of La with respect to the control input at time step ++1-1 given by is 24412 is a scaled version of the norm of the gradient of L. by 21+k- where the scaling factor depends the exponential of the real part of the Koopman eigenvalues, that is:
0h
0Lk
Vj E (.....m
dc 2+1-1
de 1
Proof. We now provide proof sketch of Theorem 3.1. As X, 14 + i. discretizing the diagonal matrix (using ZOH) and taking its k-th power gives us K e hidtss e chl Using this we can write
* 4/4+= = =
Σ
44-1 4411
Now applying the chain rule, we get derivatives of the loss with respect to of and c
09/ Ohr
OLk
halpy chd OCk
Rb
dr.
dmh
0il+*
08.+
0C*
0b44
0Lk
OLx
(k--))atp) ((-1))) 0Lk
R3
de
0ct+1-1
057
0ā
As 1e:01
1, we get the result of partial derivatives given in Theorem 3.1.
B ADDITIONAL DETAILS ON THE DERIVATIONS
We first provide additional steps for the derivation of Eq. (5). As 8141 Kx, + Lug, unroll it to get future predictions up to t time steps that is:
h++1 = Kxs + Liue
8 Ki Luy11 R(Kr, Lu,) Lu K2I, RLu, Lu4+1
.t-3 K KKK2 rxt KLu Lu+11) Luut2 R3xtt Laut R KIau4 +
Selr K KT- + RT-2 + +... Luter-1
Writing the above equations in a matrix form and using C++* Lue++ we get Eq. 5 Now to derive Eq. (7). we use the property matrix exponential of diagonal matrices. In particular, we have K (i. Am) implies that K = diag(A{..... ... Using the notations from Section 3.2 we get that for i-th index of predicted vectors #1t# the following equations hold:
# X,ay + of a a X 2 2 Aset + ct+1
#/4 X7 x + X c/41
6
The above equations can be denoted by an expression using the circular convolution operator as written in Eq.
15

---

Published as a conference paper at ICLR 2024
C PROPOSED MODIFIED DYNAMICS MODEL IN RL AND PLANNING
C.I REINFORCEMENT LEARNING
The RL setting can be formalized Markov Decision Process (MDP), where an agent learns to make decisions by interacting with an environment. The agent's goal is to maximize a cumulative reward signal, where a reward function r(s, a) assigns a scalar value to each state-action pair (s, a). The agent's 's policy, denoted by ( is), defines the probability distribution over actions at given state. The agent's 's objective is to learn a policy that maximizes the expected discounted sum of rewards over time, also known as the return: (T) Ex (E ('r(st, at)] where y E [0, 1] is the discounl factor and 8 and a are the state and action at time step I, respectively. The agent's s expected cumulative
reward can also be represented by the Value function. which is defined the expected cumulative reward starting from a given state and following a given policy.
Value-based RL algorithms have two steps of policy evaluation and policy improvement. For example, Q-learning estimates the optimal action-value function, denoted by (s, a), that represents the maximum expected cumulative reward achievable by taking action in state and improves the policy by setting it to arg max, (s, a) (Mnih et al., 2013 2015). A policy-based RL algorithm directly optimizes the policy to maximize the expected return. Policy gradient methods update the policy by adjusting the policy function's 's parameters by estimating the gradient of J. Actor-critic methods combine Q-learning and policy gradient ones, such that the Q-network (critic) learns the action value under the current policy, and the policy network (actor) tries to maximize it Haarnoja al. 2018); Lillicrap et al. 2016); [Schulman et al. (017).
C.2 FORWARD DYNAMICS MODELING IN RL
Dynamics models have been instrumental in attaining impressive results various tasks such as
Atari (Schrittwieser et al., 2020; Hafner et al. 2020) and continuous control (Hafner et al. 2020
Jiang al., 2020 Sikchi al.) 2021: Lowrey et al. 2019). By building model of environment dynamics, prodel-based RL can generate trajectories used for training the RL algorithm. This can significantly reduce the sample complexity in comparison to model-free techniques. However, in model-based RL, inaccurate long-term predictions can generate poor trajectory rollouts and lead to incorrect expected return calculations, resulting in misleading policy updates. Forward dynamics modeling has also been successfully applied to model-free RL to improve the sample efficiency of the existing model-free algorithms. These methods use it to design self supervvid auxiliary losses for representation learning using consistency in forward dynamics in the latent and the observation space (Jaderberg et al. 2017: Schwarzer et al. 2020 Srinivas et al. 2020)
C.3 BASEL-BASED PLANNING
Model Predictive Control (MPC) is control strategy that uses dynamics model s1+1 f(s,, ae) to plan for a sequence of actions a,, a.+1..... a++r-1 that maximizes the expected return over a finite horizon. The optimization problem is:
[++-1 Σ 1=1
arg max
E
Ar((s, a.)
(8)
where 7 is typically set to 1, that is, there is no discounting. Heuristic population-based methods, such as cross entropy method Rubinstein & Kroese (2004) are often used for dynamic planning.
These methods perform a local trajectory optimization problem, corresponding to the optimization of a cost function up to a certain number of time steps in the future (Williams et al., 20155: 2018
Okada & Taniguchi| 2020). In contrast to Q-learning, this is myopic and can only plan up to a certain horizon, which predetermined by the algorithm.
One can also combine MPC with RL to approximate long-term returns of trajectories that can be calculated by bootstrapping the value function of the terminal state. In particular, methods like
TD-MPC (Hansen et al.) 2022) and LOOP (?) combine value learning with planning using MPC.
The learned value function is used to bootstrap the trajectories that are used for planning using MPC. Additionally, they learn a policy using the Deep Deterministic Policy Gradient (Lillicrap et al. 2016)
or Soft Actor-Critic (SAC) (Haarnoja al. 2018) to augment the planning algorithm with good
16

---

Published as a conference paper at ICLR 2024
proposal trajectories. Alternative search method such Monte Carlo Tree search (Coulom 2006 is used for planning in discrete action spaces. The performance of the model-based planning method heavily relies on the accuracy of the learned model.
C.4 KOOPMAN TDMPC
TDMPC uses an MLP to design Task-oriented Latent Dynamica (TOLD) model, which learns to predict the latent representations of the future time steps. The learnt model can then be used for planning. To stabilize training TOLD, the weights of a target encoder network 99-, are updated with the exponential moving average of the online network go-
Since dynamics modeling is no longer the only objective, in addition to the latent consistency of Eq. (6) the latent x' is also used to predict the reward and Q-function, which takes the latent representations as input. Hence, the dynamics model learns to jointly minimize the following:
CTOLD - c1 Σ
x Il - 90- (8++ 02
>11l (o r242
C3
11Q00(00+) ((+x) -
(9)
where Rs, To. Qo are respectively reward, policy and Q prediction networks, where yt+* = Pt+k Qo- (8 (+(+++))) is the 1-step bootstrapped TD target. Moreover, the policy network is learned from the latent representation by maximizing the Q value at each time step. This additional policy network helps to provide the next action for the TD target and heuristic for planning. For details on the planning algorithm and its implementation, see Hansen et al. (0022).
C.5 KOOPMAN SELF- PREDICREE REPRESENTATIONS
To avoid overconstraining the latent space using the dynamics model, SPR uses projection heads. Similarly, we encode the Koopman observables using projection layer above the representations used for Q-learning. Using & for the input pixel space, the representation space for the Q-learning is given by 2; (s()t) where eœ is CNN. The Koopman observables I, is produced by encoding 2 using a projector po such that Hy = (z().. Moreover, it can be decoded into 2; using the decoder do.)) Loss functions the prediction loss Eq. (6) and TD-error
LSsl = Σ (do(fr+k) - ce- (5:+x)11 [[Qo az) - (re + Qo (co ($141), a4+1))l (10)
Here, 0+++1 is sampled from the policy л. The policy is learned from the representations using Soft Actor-Critic (SAC; Haarnoja et al. 2018). As opposed to SPR, we do not use moving averages for the target encoder parameters and simply stop gradients through the target encoders and denote it as eg- Moreover, we drop the consistency term in the Koopman space in Eq. (IO) as we empirically observed adding consistency both in Koopman observable (x), and Q-learning space (2) promotes collapse and makes training unstable, resulting in higher variance in the model 's performance.
D COMPARISON WITH DSSM
While our proposed Koopman model may look similar to a DSSM Gu et al. (2022a); Gupta et al. (2022); Mchta et al. (2022), there are four major differences. First, our model specifically derived for dynamics modeling with control input using Koopman theory and not for sequence or time series modeling. Our motivation
make the dynamics and gradients through time stable, whereas, for DSSMs, it is to approximate the 1D convolution kernels derived from HiPPO et al. (2020) theory. Second, a DSSM gives a way to design cell that is combined with non-linearity and layered to get a non-linear sequence model. In contrast, our Koopman-basse model shows that simple linear latent dynamics can be sufficient to model complex non-linear dynamical systems with control. Third, DSSMs never explicitly calculate the latent states and ignore the starting state. Our model works in the state space, where the starting state is crucial to backpropagate gradients through the state encoder. Fourth, a DSSM learns structured convolution kernels for ID 1D signal mapping that
17

---

Published as a conference paper at ICLR 2024
for higher dimensional imput, it has multiple latent dynamics models running under the hood, which are implemented as convolutions. In contrast to this, our model runs a single linear dynamics model for any dimensional input.
E DETAILED RESULTS FROM THE MAIN TEXT
In this section, provide the numerical values in Table L and Table 2 of the results used in Fig. Jablef]providee the numerical values for training speed used in Fig. In Fig. 7 and Fig. X] we further provide the environment-we wise returns evaluated after every 5000 training steps for the planning and RL experiments used in Fig. 5 Fig.
Table I: Forward state prediction error in Offline Reinforcement Learning environments. The reported number the mean square етrоr of the state prediction over horizon of 100 environment steps. The results include standard deviation over runs.
State Prediction Error (L) GRU
Offline Datasets
MLP
Transformer
DSSM
Koopman
expert-22
0.500 0.080
0,060 = 0.007
0,006 + 0.001
0,030 +-2 0.008
0.007 0.000
Hopper-nadi
0,550 = 0.120
0.063 = 0.009
0.011 + 0.004
0.050 + 0.008
0.012 = 0.004
kopper-full-repe
0.880 0.150
0.240 0.070
0.050
0.005
0.180 0.050
0.004
walker2d-expert-v22
0,300 0.060
0.130 + 0.060
0.100
0.005
0.140 + 0.020
0.090
0.008
walkeer2dimaddiim--eprrr---
1,320 0,210
0.310 + 0.070
0.190
0.050
0,320 +
0,220
0.040
wa wolkerr2d-fol1-rappaa ay-v2
4,500 2,100
3.300
0.350
0.850
0.080
1,200 + 0.090
0.910
0.050
halfccleetssb-exeeLL-zZ
2,860 1.100
0.840 0.160
0,280 0.050
0.520 0.050
0.300
0.060
hal balTchhhhhh--mmddu---eeeee
0.590 0.090
0.540 + 0.071
0.501 0.031
0.910 + 0.040
0.568
0.040
0,850 0.200
0.930 = 0.140
0.750 0.030
0,950 + 0.060
0,700 0.040
Table 2: Forward reward prediction error in Offline Reinforcement Learning environments. The reported number is the mean square errоr of the reward prediction over horizon of 100 environment steps. The results include standard deviation over 3 runs.
Reward Prediction Error Transformer
Offline Datasets
MI.P
GRU
DSSM
Koopman
expeer-expert-v2
0.080 0.005
0.030 :: 0.002
0.001 = 0.000
0.002 ::: 0.000
0.001 :: 0,000
0,260 = 0.050
0.046 = 0.008
0.001 : 0.000
0.005 1 0.000
0.002 + 0.000
hopper- foll-r--ull--eppaa--
0.400 0.100
0.110
0.040
0.009 0.005
0.004
0.012
wa walkerrd-expertt-v2
0.015 0.003
0.004 + 0.001
0.004 +: 0.001
0.006 + 0.001
0.003 0.000
0.190
: 0.220 = 0.020
0.011
walkeer2dd-full--rpppaa-zz
0.520 0.080
0.080 = 0.015
0.091 = 0.005
0.090 + 0.008
halfchhettsb-exxeitt-v
0.160 0.050
0.070 0.005
0.045 0.004
0.120 0.020
0.040 0.006
0.230 0.050
0.120 = 0.070
0.050 = 0.007
0,150 = 0.007
0.060 + 0.005
salfchhetaa--tull-eeeeel
0.170 = 0.045
0.110 - 0.052
0.105 + 0.009
0.205 1 0.011
0.091 +: 0.006
Table Training speed (iterations/second)) for a state prediction task using an MLP, Transformer, and Kiggonal-Koopman dynamics model halfcheetalh -eppeeetteeeeeeeeeeeeeeeeeeeeeee v2. Each iteration consists of one gradient update of the entire model.
Training sequence length
Dynamics Model Type
50
100
200
500
MLP
478.00
115.54
55.85
28.02
11.20
Transformer
418.00
120.20
73.34
34.48
14,92
GRU
420.00
101.20
42.02
17.40
8.91
DSSM
530.32
148.50
76.10
37.12
17.20
Koopman
565.43
160.76
86.49
59.65
24.56
18

---

Published as a conference paper at ICLR 2024
Acrebot Swingup
Cheetah Run
800
Episado raturn
400
200
malr iopprsss soor Ayoorics
ToeRt Ineprsss Aywrris
0.2
Quadruped
Walk
1000
800
600
600
400
200
TDNPCC dyvaric venila TONAC
TCPECC dynrria velila TDMPC
0.0
0.1
0.2
0.3
0.5
0.8
0.3
0.3
0.5
Environment stops (x10f)
Environment steps (x105)
Figure 7: Comparison of our Koopman- based dynamics model (with a horizon of 20) with an MLP- based dynamics model of vanilla TD-MPC (DD-MPC[Hansen al. (2022). The results shown with 5 random seeds with the standard deviation shown using the shaded regions.
19

---

Published as a conference paper at ICLR 2024
1000
800
330
800
700
160
230
340
400
400
100
300
200
kougran-bwwed dynumics
kosprattbbaee dynam
loopman--aaee dymtmmss
rip based dynamics
200
tp-osed dynamics Vanilia AL
mpp based dynomics
Vanilla
varilla RL
0.0
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.3
0.2
0.3
0.4
05
Instronment stwps 1x30%)
Imainnment steps 1x
Environment steps stepe (x
Spe
Walker
600
500
600
400
Splasd
lossrann-aaee dynamics based
coopraan based dyram baseeadd dynamics
022
Figure 8: Comparison of vanilla RL (SAC Haarnoja et al. 2018) and its integration with MLP- based (SPR Schwarzer et al.) (2020)) and a Koopman self-predictive representations the DeepMind Control Suite. Shaded regions represent the standard deviation.
Koopman-based based dynamics model for incorporating
F MORE DETAILS ON THE BASELINE DYNAMICS MODELS
For all the experiments, we used a model with around 500k trainable parameters to ensure that the improved compute efficiency during training is not due to the model having fewer parameters. To incorporate causal transformer architecture with that many parameters, we reduced the embedding dimension of both the state and action embedding to 64. We further reduced the number of heads to 4 to design a causal Transformer model that matches the size of our Koopman model. For the GRU as the Koopman
and MI.P-based model, keep these dimensions the dimension size 64 for the GRU based-model in order to keep the parameters similar. Moreover, all our experiments in Table 3 use one Nvidia A100 GPU for training.
We use hidden
G ABLATION STUDY ON INITIALIZATION
We now perform ablation study on the different possible initialization methods presented in Section 3.2.2 We report the results for halffchetthh-expert--ee walker-2d-expert- and hopper----e.oo....... expert-v2, in Table [ndd We see constant value of -0.2 for μ4 and increasing the order of frequency for w; gives a consistent performance boost over different environments.
20

---

Published as a conference paper at ICLR 2024
Table 4: Ablation on initialization of the eigenvalues of the diagonal Koopman matrix: A, pu + jws for hal wcceeteah-epeeeeeeee.. The results are over three seeds.
State Prediction Error (4) Reward Prediction Error (J)
constant
increasing
0.300 + 0.060
0.040 + 0.006
constant
random
0.340 = 0.047
0.041 + 0.005
learnable
increasing
0.320 = 0.024
0.048 1 0.003 0.042 + 0.007
learnable
random
0.360 + 0.040
Table 5: Ablation on initialization of the eigenvalues of the diagonal Koopman matrix: X = + Jww for ailker-2d-expert The results are over three seeds.
f4
W/
State Prediction Error (L) Reward Prediction Error (L)
constant
increasing
0.090 + 0.008
0.003 1 0.000
constant random
0.104 = 0.008
0.004 + 0.001
learnable
increasing
0.102 0.003
0.004 + 0.000
learnable
random
0.096 + 0.007
0.003 + 0.002
Table 6: Ablation on initialization of the eigenvalues of the diagonal koopman matrix: A; 14 + jw: for hopper-expert v2. The results are over three seeds.
14g
wi
State Prediction Error (L)
Reward Prediction Error (L)
constant constant
increasing random
0.007 + 0.000
0.001 0.000
0.010 + 0.007
0.014 + 0.005
learnable learnable
increasing
0.008 + 0.006
0,002 + 0.004
random
0.008 + 0.004
0.003 0.001
H
ABLATION STUDY ON HORIZON FOR PLANNING AND MODEL-FREE RL
We investigate the role of horizon (or rollout length) in both model-based planning (Section 4.2. and model-free RL (Section 4.2.2) for our Koopman-based dynamics model. The corresponding results are available in Fig. 9 and Fig. TO, We observe that for model-based planning, the performance improves with longer horizon. However, for model-free RI, the performance is higher with prediction horizon of and deteriorates for 20. This phenomenon consistent with the results
reported in Figure 4 of Schwarzer et al. (2020).
Mean
IQM
Optimality Gap
(horizon
Koopman (horlzon 20)
0.750 0.765
0.780
0.785
0.675 0.690 0.705 0.720 0.735 0.270 0.285 0.300 0.315 0.330 Human Normalized Score
Figure 9: Ablations of our Koopman based dynamics model for horizon lengths of 5 and 20. The results are over 5 random seeds for each environment. Higher Mean IQM and lower Optimality Gap is better.
21

---

Published as a conference paper at ICLR 2024
Mean
IQM
Optimality
Knopman thorzon 20)
Koopman (horizon 5)
0.545 0.660 0.675 0.590 0.705
0.630
0.660
0.345 0.375
Human Normalized Score
Figure 10: Ablations of our Koopman based dynamics model for horizon lengths of 5 and 20 for incorporating self-pre seffprredictive representations with vanilla SaC[Haarojaa et al. (2018) in the DeepMind Control Suite. The results are over random seeds for each environment. Higher Mean IQM and lower Optimality Gap is better.
I ABLATION ON THE DIMENSIONALITY OF LINEAR LATENT SPACE
We also ablate the dimensionality of the Koopman state observables: m in Table 7 and show that 512 is enough for most of the environments, the performance saturates there.
Table Forward state prediction error Offline Reinforcement Learning environments. The reported over
number is the mean square error of the state prediction in hal horizon 100 environment steps. We observe a general trend a decrease in the error values with an increase in the state embedding dimension which is aligned with the Koopman theory. We have selected an embedding dimension of 512 for all our experiments as has the lowest error value.
State Embedding Dim. (m)
State Prediction Error (J)
128
0.0097
256
0.0081
512
0.0073
1024
0.0075
J KOOPMAN DYNAMICS MODEL IN JAX
We provide the Jax code snippet efficiently implement Koopman-baa bbased dynamics model. Our code can be found in http::/ggthh..ooommmmmm///ooonnn
22

---

Published as a conference paper at ICLR 2024
Algorithm 1 Diagonal Koopman Dynamics model
Enport jax
impoxt jax-numpy as Top
from www.nn.iittiiiieee. foom flax import
nornzl
Kain class to implament Nisrcnal Xocenan Dynanics modal
class srcccaiiisphalllppppppanntttno (an, dondeeell))
int
int real_init_ sex 'constant'
reol_init_vville value: flest -0.3
im init type: str linrressing freq'
istivaaiooss Callable (allable[[] [jdp.nderryy,, jop.nderray] mn.ell
koognen ioogner.nnnginaar inw.inary. (selfis
self. inpe "increasi no..iing...ooo... freg"s
selr, in self.parant "Kinn" increasing_in
(selt.atata_din, // 4,3
elif self,in
'randon';
seir, self-paran(
"Kii"" tandon (sclt.stata_din 4,3 else:
ValueErrGE
def init_ koognon roaman__roa__n pasl_sssssttt))
if self, self.reel... reel__nit__yp type =- 'constant's
falf,8 real self, ieal olif self.reai_init_ee
'lcarnable':
relf.K_= pelt.parant "K real",
(sel 4,1
self.real cerreea__iiiittteee vnlt_velue
selr I.p.c...aaii.RRR -0.4, -0.1)
elset
ValueFrror
daf (solf):
model porameters
self. init weapt__oooeeneaelllcoooe
solf.init_koonne
sclf.action.ehhcdee MLP (128, self.acticn_mmb.dtt 2],
self.1 celt.paran(
(0.1), (seli.action_enb din, self.state_din 23
Step paremater
self.step jnp.exp ("log log.oooooooooo step", log_step I1 1,11
xompxxo jeg.concaterater
[roll - 11 2012.K_in,
welf,Rrrea 1j sel self,Kinn),
axis-
solf.L_compiox sele. 5 +1 solf.ifi, & scif.x_dis, seif.lddie discretizel seippR_ 2elf. iooppeen self.atep
def cell (self, accions, start statt_stttee_re state_rep);
12 (sctionns shapel seturn caaprn___orrrr_cg__oooo
self, cis,, self-lidis,
(act ons), start__tate__reg
else:
Vand inp.vander
&.r.R&Giis act iors.shapee1] IncreasingeTrue
foopman_.orxaadd
Vand_R, self fi
self. ael.. iondeccdee (actional, atart_state_ree
23

---

Published as a conference paper at ICLR 2024
Algorithm Helper functions for Diagonal Koopman Dynamics model
causal_ccane cow.oooooooooooooooo
assext shape s.spape
jap. Eft. (je,-pad IC, Shahape:0011 (01:))
jap.ctt. fft. (Jnp-pad (K, 16, s.shape out ud Kdl
itg.ftt.ifft
a.ehape (01]
computo ceapett_meecuccmnnttooo www.eurconntt.eeeeeeeeoooooo actior_enb, inis_statel::
return caussl coneol_oovollrrioiiee wact:rooooiaiiiooeocc.o. xer(i-1]1 Rec]l:) itat_state
def mesaurement (Ker, actloon init_atate):
jax.vrap!
compute meneureentt block,
4xee=(0, 1, 01,
pas_ases=] (Bez, action_end, init_atate)
discretize
atepl:
return jag-exp (steprit, exp-exxttteepx3 -1/K
det jog_stop_intis init (dt (dt__in=0.0000 001, dt_nexm0.1::
det init (key, shape):
www.:aann..iifrrnnooommmm (key,
(dx___ma) jap.lor(d=_min) logit_og
return inst
dst increasing_in_iaee
init ikey, shape):
return jap-pi (op.arange(shage(0))
seturn init
def randon in in
def init shape):
zetera jax. wnd.ddd....ffftnjkke... slape)
return init
ejit coopan__orward/aadd....
actions_enb, start__tatterop)
astion_ent_dia aetions_emb.chapcls
state_dim dim stss___sttt_rrppspapptt
actions cicions_erb_copplee acticns embli i, saction enb aen_neeenndd/daaeed /2] \
+ +f actiens_anblt. astion_ene, atteoo______nna//a_eeeeeee din//2:1
start_etatc_rep_comple
start__tatt_repl: istate_din
+i start_atase_repll state_din
predicted staie, rem____ep_mmmmeee jax. cemputo__marureneen in axes- (None, 0, 0),
[Wad_ K, cetiona_embb_comple. i, slart__iiall__ep__ooobaa.
predicted seciicted__tttt state rep jop, concateneeee jap, seal (predicted pttp imag bowp:/dottoooattttooeeeeeeeeeeee
return prodictoi_statc_rop
81it
soopan_.orword_sinngo
acticn_cab, start_stats_res):
action aation_enb_dien _nn__din
state dim start state rep.atapel-1
setioo_ent_ccmplxx action_entile It, sacticn_onh_ditr// cin///2 1
- action_enblis action_enb. 1isoo_omb_ddm//211
stare seate_rep__cmpeer rep_complex
din
start_state repls, state_din
ppadicte__tatterrppoommpl saction _ complo 8
K[None,
predicted state rap jap. concesteeeee
jap. seal (predicted aitee
jap. imag oeggg
return predicted_atate...
24


Reviews:
[{"review_id": "IrybGZy95M", "reviewer": "Reviewer_xZWp", "summary": "The paper proposes an approach based on Koopman theory, which linearizes the nonlinear dynamics of the environment in a high-dimensional latent space for long term prediction. This allows for efficient parallelization of the sequential problem of long-range prediction using convolution, while considering the agent's actions at each time step. The approach also enables stability analysis and better control over gradients through time, resulting in significant improvements in efficiency and accuracy of modeling dynamics over extended horizons.", "questions": "1) It would be interesting to see the results with NeuralODE\n\n2) Please do a more thorough literature review on Koopman operator theory for control and modelling.", "limitations": "", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The strengths are as follows:\n1) They provide extensive experimental results verifying the claims made in the better and showing better long term prediction accuracy that other recent non-Koopman NN based methods.\n\n2) They provide numerical simulations that they Koopman based approach does not lead of instability. In other words, they show that their long term predictions do not blow up.", "weaknesses": "The weakness of the paper are as follows:\n\n1) I've had the opportunity to delve into Koopman operator theory in my past research. It's noted that control-affine systems transform to a bilinear control Koopman based system in the lifted space under particular conditions, as exemplified by Theorem II.1 in the paper titled \"Advantages of Bilinear Koopman Realizations for the Modeling and Control of Systems with Unknown Dynamics\". It's also worth mentioning that these bilinear forms may not be universally applicable to all nonlinear systems, especially those that do not adopt the control affine form. Some recent works on Koopman operator theory have also showcased how a control-affine nonlinear system can be transitioned to a more generalized input-separable Koopman system, with bilinear and linear forms being special instances of these separable Koopman formats.\n\n2) There are alternative NN methods, such as Neural ODE, that have consistently demonstrated promising results for long horizon predictions, in my experience. It might be valuable to consider or reference them in the context of this paper.\n\n3) The problem of long horizon prediction for both controlled and non-controlled dynamical systems using Koopman operator has been explored before, please see the Nature publication titled \"Deep learning for universal linear embeddings of nonlinear dynamics\", and in subsequent studies.  I am aware that you already cited this reference but “twice” (please check)\n\n4) While controlled systems are a primary focus in this paper, an additional section dedicated to control design could potentially improve the paper further.\n\n5) On the topic of stability in Koopman-based learned matrices, there are several contributions, such as \"Learning Stable Koopman Embeddings\", and its subsequent related studies. It might be advantageous to discuss or incorporate the stability guarantees they present in the context of this paper.\n\n6) Considering the above, a more comprehensive literature review on the Koopman operator for controlled/non-controlled dynamical systems might enhance the paper's breadth.\n\n7) For the purpose of replication and validation, making the code available as open-source (while preserving the authors' anonymity) could be beneficial for the broader community.", "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": "", "extra_scores": {}}, {"review_id": "xeRW6XZh9c", "reviewer": "Reviewer_UaZL", "summary": "In this paper, the authors introduce a diagonal Koopman-operator approach that helps to efficiently learn and quickly plan with dynamic models. The standard encoder used in Koopman-operator-based approaches is modified to output the complex eigenvectors of the latent linear model, whose estimation is then reduced to learning a diagonal matrix (together with the encoder/decoder weights). By careful initialization of the complex eigenvalues of the Koopman-matrix the authors can also protect against vanishing/exploding gradients. The diagonalization helps quite significantly to reduce the computation time needed for planning over long horizons, as the authors show in several dynamics-estimation and RL based tasks. The method also works competitively with several state-of-the-art approaches in dynamics modeling and RL applications.", "questions": "* This is not a problem only concerning this work but in general applies to all 'Koopman-operator-inspired-approaches': the connection to the theory is very very loose and hence these methods could instead be simply called linear latent-dynamics modelling. Do you agree? If not, please mention how the Koopman-operator theory guides you analytically and intuitively. \n\n* As a follow-up to the above: how does the encoder quality effect your results? How do you come up with the encoder structure? Can the theory guide you here?\n\n* In the optimized Koopman matrix, are the eigenvalues sparse, or are they all nonzero/complex ? Can they be made sparse and is that interesting? \n\n* Please mention the limitation to deterministic environments in the introduction. Is this a limitation that all Koopman-operator approaches in the literature face?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The approach as far as I can tell is sound, and as the summary above indicates, the authors show significant improvement in terms of computation time / backpropagation stability that results from the diagonalization and the careful initialization. The method is also tested in several benchmark problems (RL and dynamics modeling) and the ablation studies in the appendix look comprehensive to me.", "weaknesses": "Sometimes the discussions in the paper are not very clear, for instance\n\n* \"While the same dynamics model can also be used for (III) model-based RL, we leave that direction for future work.\"\n> Do you mean you would improve the model K further? or g_theta perhaps?\n\n* Can your approach fail even if eigenvalues are initialized well? After correct initialization, will the real part of the eigenvalues always stay negative? Are there any other failure scenarios?\n\n* \"we apply our dynamics model in model-free RL, \"\n> What does it mean to apply a model in the model-free setting? Please explain not just in the appendix but also in the main text briefly. In general much of the discussion in the Appendix could be streamlined with the text. Please organize the appendix and reduce its scope, code should be put in github and linked, no need for printing your code in the appendix. Instead you could mention more failure cases as part of the ablation studies perhaps.\n\n* Would be nice to have figures or a section explaining *what* you're solving, rather than how, e.g. inputs are in the pixel space, encoder is an MLP and the TD-MPC tries to solve X ... So some part of Section E from the appendix could go to the main text.\n\n* It is not clear to me which network architecture you used for the encoder g_{\\theta}, is it another MLP?", "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": "", "extra_scores": {}}, {"review_id": "AgryCE2EQy", "reviewer": "Reviewer_3syb", "summary": "This paper introduces a new method for learning a Koopman embedding that gives decoupled linear dynamics in a latent space, opposed to nonlinear dynamics in the state space. This approach is also able to specify the eigenvalue spread of the resulting linear dynamics, stabilizing gradient evaluations through long chains.", "questions": "N/A.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. There are many papers that propose using neural networks to find a Koopman embedding, this one is unique (to my knowledge) because it focuses on stabilizing the eigenvalues of the latent-space dynamics. \n2. The included JAX code is very helpful for comparing the implementation to algorithm for a better understanding of what's going on. \n3. The Koopman background section is clear and does a good job of defining all of the notation. \n4. The argument for why a Koopman model is easier to parallelize over RNN's like GRU and LSTM is a great addition to the algorithm.  \n5. This approach was shown on a wide variety of robotics tasks from the DeepMind control suite (Ball in cup catch, cartpole, swingup, cheetah run, finger spin, and walker walk).", "weaknesses": "1. Figure 3 seems to undermine the utility of the presented Koopman approach, since the GRU is able to compete with it on all of the examples and metrics. Maybe some explanation in the figure caption explaining why this is expected would be appreciated. \n2. Longer/better figure captions would be nice. \n3. This seems like a marginal improvement over existing works cited. Some more comparisons between this approach and existing algorithms could better showcase the additions.", "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": "", "extra_scores": {}}, {"review_id": "7Qpk40B41K", "reviewer": "Reviewer_tVSX", "summary": "The authors propose a new state-space model for dynamical systems modelling based on the Koopman operator. The authors approximate the Koopman operator in such a way that they can leverage time convolution operators and compute multiple-step predictions without resorting to autoregressive filtering. This seems to reduce training time and improve learning algorithm stability. This is because the gradient growth can be directly controlled due to the structure of the proposed model. The authors also show that replacing a native model in some model-based algorithms with the proposed one improves the algorithm's performance.", "questions": "1. “In this work, we leverage techniques and perspectives from Koopman theory (Koopman, 1931; Koopman & Neumann, 1932; Brunton et al., 2021) to address this key problem in long-range dynamics modeling of interactive environments.” I feel that Mauroy et al 2020 would be a better reference for modern Koopman theory and recent developments \n1. “and thus an invariant subspace $G \\subset F$ is often used, so that $Kg \\in G$. $G$ is spanned by a finite set of observables $g_1$, . . . , $g_m$, where often we assume $m \\gg n$.”\\\nThis is slightly confusing. Finding an invariant subspace *is not a tractable problem*, while this phrasing implies that we can pick an invariant subspace. Please rephrase.  \n1. Page 3. “one could show that the Koopman operator is bilinearized (Brunton et al., 2021)”\nIf I understood correctly Brunton et al bilinearization is possible only if there is a finite number of eingenfunctions (Theorem 6.1). I am not sure if this implies isomorphism to linear dynamics, but in any case, this is an important restriction worth mentioning \n\n1. “An alternate approach to simplifying the Koopman operator with a control signal assumes”. “Simplifying” is a confusing word in this context. I recommend using “approximating”, which would rule out equivalence of representations.\n\n1. “Fig. 4 empirically verifies that our proposed Koopman dynamics model is around twice as fast as an MLP, GRU or Transformer based dynamics model in learning from longer trajectories” It is hard to verify this claim based on the figure, I suggest referencing Table 3 in the appendix and provide the number 1.7 instead of “around twice as fast”\n\n1. There seems to be a small clash of notation, where $\\overline \\lambda_j$ is used to describe the entries of the diagonal Koopman matrix $\\overline K$, and also the diagonal entries of the matrix $K$. \n\nReferences:\n \nJanner, Michael, et al. \"When to trust your model: Model-based policy optimization.\" Advances in neural information processing systems 32 (2019).\n\nFujimoto, Scott, and Shixiang Shane Gu. \"A minimalist approach to offline reinforcement learning.\" Advances in neural information processing systems 34 (2021): 20132-20145.\n\nMauroy, Alexandre, Y. Susuki, and I. Mezić. Koopman operator in systems and control. Berlin: Springer International Publishing, 2020.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea of creating a linear approximation of the nonlinear dynamical system through a Koopman framework is very interesting. This is because we can use a well-developed linear systems theory to derive scalable algorithms. \n1. The theoretical result on gradient growth is valuable for practitioners, who are not well acquainted with dynamical systems theory and stability theory\n1. The authors demonstrate on some numerical examples, that the generative model is efficient for learning a planning controller as well as a policy.", "weaknesses": "1. Contributions are not well articulated in the introduction. It is not clear how this work is different from other Koopman operator-based RL methods. For example, I find the tools to compute long-term predictions quite interesting and a bit more detail on this in the introduction would strengthen the paper.\n1. Ablation on sequence length for RL tasks is not provided, i.e., do we get an improvement from the model being able to predict several steps ahead? This is an important ablation since it was shown by Janner et al that using SAC as a policy learning algorithm does not benefit from rollouts longer than a one-time step. One potential reason for this is the algorithm itself that learns only from one-step-ahead predictions. Therefore, this raises the question of whether we actually need a model capable of long-term predictions. I do not see a discussion on this subject. \n1. Ablation with Hippo model as a backbone for offline RL is missing. The authors mention the relation to the Koopman framework and it would be interesting to see if the proposed model offers an improvement over HIPPO\n1. The presentation needs to be sharpened a bit to highlight the main contribution - the effect of the novel Koopman model on RL in comparison to other generative models. I think the authors can pivot to argue that their Koopman model should be considered as the base model for most of the RL problems (offline RL, planning, model-based RL etc). But this a matter of choice.\n\nI will consider my score further in the review process", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)