Course Correcting Koopman Representations
Authors: Mahan Fathi
Clement Gehring
Jonathan Pilault
David Kanaa
Pierre-Luc Bacon
Ross Goroshin
Venue: ICLR 2024 poster (2024)
Keywords: Koopman
Autoencoders
Dynamical Systems
Sequence Modeling
Inference-time Methods
Planning
Unsupervised Learning
Representation Learning
Robotics
Primary area: applications to robotics, autonomy, planning
TL;DR: We use Periodic Reencoding as means for correcting the drift of Koopman representations at inference time.

Abstract: Koopman representations aim to learn features of nonlinear dynamical systems (NLDS) which lead to linear dynamics in the latent space. Theoretically, such features can be used to simplify many problems in modeling and control of NLDS. In this work we study autoencoder formulations of this problem, and different ways they can be used to model dynamics, specifically for future state prediction over long horizons. We discover several limitations of predicting future states in the latent space and propose an inference-time mechanism, which we refer to as Periodic Reencoding, for faithfully capturing long term dynamics. We justify this method both analytically and empirically via experiments in low and high dimensional NLDS.

COURSE CORRECTING KOOPMAN REPRESENTATIONS
Mahan Fathi
Clement Gehring
Jonathan Pilault
Google DeepMind, Mila Université de Montréal
Mila
Mila
Université de Montréal
Polytechnique Montréal
David Kanaa Mila
Pierre-Luc Bacon
Ross Goroshin
CIFAR AI Chair, Mila Université de Montréal
Google DeepMind
ABSTRACT
Koopman representations aim to learn features of nonlinear dynamical systems (NLDS) which lead linear dynamics in the latent space. Theoretically, such features can be used to simplify many problems in modeling and control of NLDS. In this work we study autoencoder formulations of this problem, and different ways they can be used to model dynamics, specifically for future state prediction over long horizons. We discover several limitations of predicting future states in the latent space and propose inference- time mechanism, which refer to as Periodic Reencoding, for faithfully capturing long term dynamics. We justify this method both analytically and empirically via experiments in low and high dimensional NLDS.
INTRODUCTION
Recent research has shown a growing interest in learning representations of nonlinear dynamical sys- tems (NLDS) in which the dynamics become linear in the latent space. Linear dynamics in the latent space offer distinct advantages. including the capability to derive closed form solutions to optimal control problems using LQR solvers (Kalman, 1960). System identification and interpretability are also greatly simplified for linear dynamical systems. From computational standpoint, advancing linear system forward can be executed more efficiently by leveraging parallelism (Gu et al., 2022; Smith et al., 2023). An example of this approach is presented in S4 by Gu et al. (2022), carefully designed deep state-space model (SSM) that takes advantage of the parallelism offered by discrete linear dynamics.
Given the advantages of linear systems, a challenge in the study of dynamics and data-driven modelling resides in extracting global linear representations of nonlinear systems. In this regard, Koopman Theory (Koopman, 1931; Koopman & v. Neumann, 1932) provides a framework in which nonlinear dynamics can be cast into linear ones in a space of measurements, spanned by basis of characteristic functions, which are the cigenfunctions of Koopman's 's composition operator. However, identifying such characteristic functions as well as the structure the linear transition operator can prove very difficult, in general, as they depend heavily on the qualitative dynamical properties of the system under study, necessitating the need to resort approximate methods.
Recently, several studies have explored the integration of deep learning architectures with Koopman operator theory, primarily in the context of small scale well-behaved problems. Azencot al. (2020) train a nackward-compatioe by assuming reversible dynamics, ensuring stability during training by driving the eigenvalues of the operators to be close to one. Frion al. (2023) adopt similar strategy by promoting orthogonality of the Koopman matrix through an auxiliary loss. Lusch et al. (2018a) Iearn generalized Koopman representations for systems including those exhibiting continuous spectra. Mondal al. (2023) employ a Koopman- bassed model in the context of Model- Reinforcement Learning as a short- horizon (20-step) planner.
Koopman matrix in conjunction with its original forward counterpart,
In this work, we further explore the use of principles outlined by Koopman theory (Brunton et al., 2022), guiding framework to obtain linear representations of NLDS within autoencoder

---

framework. We observe that unrolling the linear dynamies in latent space leads to long term drift in trajectories when mapped back to the space of observables (state space). Prior works achieved stable unrolls at inference time by restricting the cigenvalues of the Koopman/linear operator, and training on long sequences. We posit that such requirements, i.e. long training sequences, are unnecessary, especially when modelling fully observable dynamical system. In this work uncover two limitations of generating trajectories in the latent space: (i) long horizon trajectories cross, violating uniqueness of solutions of dynamical systems, and (ii) latent space trajectory generation is unable capture switching dynamics between fixed points (Lan & Mezić, 2013). To overcome these issues, we introduce a simple inference method called Periodic Reencoding that produces high accuracy predictions over long horizons. We derive analytic expressions and present special case example (Appendix D) that provides intuition about why this method better captures long-term dynamics. Finally, we empirically validate our approach on a set of established NLDSs, and also demonstrate these findings on offline reinforcement learning datasets in more complex environments within D4RL (Fu et al., 2020).
DEEP KOOPMAN AUTOENCODERS
Given nonlinear dynamical system, Koopman theory attempts to approximate or "explain" nonlinear transitions with a linear dynamical system refer to Appendix B for a detailed overview of the Koopman theory. For example, Dynamic Mode Decomposition (DMD) (Brunton al., 2022) tries to approximate the infinite dimensional Koopman operator by fitting discrete transition data collected from nonlinear dynamical system. More specifically, using an integration scheme in time a discrete nonlinear dynamical system can be obtained from its continuous time counterpart, Ft-1 = F(x:). DMD then solves the following optimization problem over the dataset of all transition pairs, D {(1,, #c+11
2t, T++1 E J:
K DMD = arg min Σ ||+11 - Keell2.
(I)
where the K DMD E Rdxd matrix is the linit-dimensionn approximation to the linear Koopman operator. Thus DMD treats the state itself as the features, by simply fitting the Koopman matrix directly. Extensions to this approach include "extended" DMD, eDMD (Williams et al., 2015; 2014; Schütte et al., 2016), which augments the state with fixed nonlinear transformations of the state. In the spirit of end-to-end learning, are interested in data driven approaches that learn nonlinear transformations of the state (Brunton et al., 2022: Lusch al., 2018b). Koopman features can be learned in autoencoder setting by minimizing:
min Σ - + [10()+1)) Ro(x:)l K.
where the encoder Φ Rd + R", decoder 3 R² + Rd are parameterized function approximators, and À > 0 scalar. The feature vector 2, E R", representing the Koopman embedding for which the dynamics are linear, obeys the following relations:
2t ((x)), 2441 Ket.
24 T (2))
These relations are approximate because, for example, reconstructed states are approximations of the corresponding true states. Once trained, the encoder, decoder, and K E Raxx comprise a complete model of the dynamical system, potentially capable of performing long-range state prediction over arbitrarily long time horizons.
Though other feature learning methods, such contrastive learning, possible (Lyu et al., 2023), we study the autoencoder formulation because it is more pervasive in the literature. Furthermore, the learned decoders from this formulation allow for solving the control problems in latent space and mapping control signals back to phase space. This is non- conttiivial using contrastive learning which only trains encoder.
2

---

3 METHOD
3.1 TRAINING SEQUENCE
In this section formally outline the training objective for sequential data. We start by using the continuous parameterization of the Koopman dynamics, for controlled systems. An autonomous controlled system described by & = -(x,,) where is an exogenous control input. Assuming that a bounded Koopman matrix, K E IR"x", can be approximated for measurements <(x) E R", the linear dynamics then prescribed by:
d ग. [o(x) K Kő(x) + Lw(u),
where L R represents the controlled latent dynamics for external coded input v == w(u) R We optimize the objective by taking gradient steps directly over the continuous parameterization of the Koopman dynamics, i.e. K and L, which we discretize via the bilinear method (Tustin, 1947). over timestep of d:
E++1 = Kāt + Lo
(5)
where K k- 6
-
and L - 8L.
(6)
We treat trainable variable as well, and assume that samples are uniformly distributed in time. The Koopman Autoencoder architecture consists of the following trainable components, (i) the latent dynamics K, L and å (ii) the state encoder Φ (iii) the action encoder w (iv) and the state decoder vb. The training data consists of an initial state, Zt, and sequence of following actions and states, ie. (ue, T2+1, U4+7-14 I&+r) ) of length T. The model takes in the initial state and the sequence of actions input and is tasked predict the sequence of future states. To respect the Koopman dynamics, i.e. ensuring that and L are the only means for advancing the dynamics, we minimize the "Aligment", "Reconstruction" and "Prediction" losses that prevent trivial solutions.
LAlign El ||ā 6(a4+>ll2 Leccoma Ei=o || v((==))llz Prea Σ |ke+: -(A4+=l2 (7)
In the above equations, & denotes the resultant latent code after one or more applications of Koopman dynamics, while 2; represents the resultant latent state immediately after encoding (see Figure 1).
3.2 TRAJECTORY GENERATION
prediction loss
Figure 1: Course Correcting Koopman Autoen- coder Unrolls via Periodic Reencoding. Unrolls are generated by first encoding the initial condi- tion, then linearly advancing the dynamics in latent space, and finally decoding trajectory back the original space. Values enclosed within squares represent ground truth, while those enclosed within circles are inferred by the model. The objective consists of an alignment loss (shown in green). reconstruction loss (in yellow). prediction loss (in blue). We propose Periodic Reencoding: at ev- ery k steps in the figure. The latent state, 4 is decoded and subsequently reenended. Because the encoder/dceoder are not exact inverses, the reen- coded feature vector different from the original. Control inputs are omitted for simplicity.
iP--3
Fesk
generated trajectory
recncoding step
aligoaent
Pack
Iateut squeen
V
Pe-s
Peik
2212
reconstraetion
The quality of a model of a dynamical system can be assessed by generating trajectories in the state space of the original dynamical system, starting from some initial condition. There are methods for generating trajectories.
Without Reencoding, This method generates the entire trajectory in the latent space and only uses the encoder to obtain the initial condition. The decoder á is applied to every point on generated trajectory to obtain the corresponding curve in state space. More specifically, given:
i K =0==((0)) and
(>),
3

---

can write the explicit solution in state space
x(t) - ole (zn)).
(8)
It is the solution to linear dynamical system, given by the matrix exponential, in latent space which is then mapped to the original state space using tb. Under this setting, the relationship between 2- space and space is established by mapping trajectories generated in the 2- space via linear dynamics, to curves in the x-space.
In the standard autoencoder setup, the mappings between the original and latent spaces are not strictly one-to-one. This characteristic becomes especially significant when the dimensionality of the latent space is substantially larger than that of the original space, i.c. n > d, as is typical in Koopman autoencoders the inverse function theorem requires that n = d. Because of this, the curves generated without reencoding could potentially intersect with themselves, and therefore won't t faithfuly capture the characteristics of a trajectory, generated by dynamical system. This directly arises from the existence of multiple points in z- spac that lie on the same trajectory generated by the linear system, that are mapped to the same point in s-space, due to lack of injectivity in the decoder, t.. This is clearly illustrated in Figure (a) where phase lines intersect. Intersecting phase lines are impossible to generate with continuous dynamical systems because they imply that multiple solutions exist corresponding to the unique initial condition given by the cross-over point.
With Reencoding, This method uses the encoder/decoder to iteratively generate trajectory in state-space. The encoder, Koopman operator, and decoder effectively define a dynamical system in the state space, x. More specifically, given:
1= == 2 == =<),, and x z= (2),
we can express the dynamical system that implicitly defines the solution in state space as:
Ju (o(x)) K (z(z)
Where Jos is the Jacobian of v. When generating trajectories without reencoding, linearity in the latent space is assumed to hold for all times, i.e. globally. In contrast, when generating trajectories with reencoding, the linearity property is assumed to hold only locally. Equation gives rise to a dynamical system with feedback, thus the relationship between 2-space and z-space is governed by point-wise mapping of dynamics, Appendix A shows extensions to any latent state dynamics.
In theory, the trajectories generated using this method can faithfully capture the dynamics and will not produce invalid trajectories that cross, however, reencoding at every step poses two significant drawbacks: (i) repeated applications of the encoder can potentially accumulate errors much faster than repeated applications of K, as seen in Figure 4 (b), and (ii) it is not computationally efficient because unrolling must be performed sequentially due to presence of nonlinear operations at every step, unlike the parallelizability of linear operations.
We introduce Periodic Reencoding, technique spanning the middle ground between trajectory generations with and without reencoding. When generating trajectories with periodic reencoding, we decode and reencode periodically, in continuous time at A/ intervals (Figure 2), and in discrete time every ie steps (Figure 1), treating At or k as hyper-parameters. Given that the encoder and decoder are not perfect inverses of one another, the output of the reencoding step is going to be different from the original point. In other terms, the mapping between the original and latent space lacks bijectivity.
By Periodic Reencoding, we expand and harness the applicability of local linear Koopman dynamics around the initial state all the while mitigating the accumulation of encoding errors. We observe that periodically applying reencoding is an effective way to mitigate the drift accumulated by generating trajectories in the latent space. We term this property "course correction." " Using our method, we are able to generate stable, plausible, and accurate unrolls over extended horizons, while remaining computationally efficient. In Section 4, we empirically establish the effectiveness of periodic reencoding for stable and accurate long-range predictions, in contrast to unrolls without reencoding or reencoding at every step. Furthermore, we demonstrate that periodic reencoding can be beneficially integrated into the training process to achieve further improvements. Lan & Mezić (2013) showed that NLDS with multiple fixed points can only be linearized within the basin of attraction of each fixed point. This reveals another limitation of unrolls without reencoding they cannot capture the switching dynamics between multiple fixed points (see the example in Appendix D). is important to make the distinction between reencoding and "teacher forcing" (Bengio et al., 2015). Teacher forcing periodically uses ground truth during training to avoid error accumulation. Reencoding never uses ground truth data.
4

---

82 - =0:0
20 6(x0)
21 ckat 20
E3 = ckai22
REENCODE
REENCODE
2 e c
σ = (
to
t = to + At
1 = to + 2
13 = to + 3At
Figure 2: Periodic Reencoding in continuous time for trajectories in latent space. In this figure, the unrolled latent before and after reencoding are denoted as 2, and &, respectively.
4 RESULTS
We empirically evaluated our proposed approach number of highly nonlinear environments with varying dimensionality. begin by modelling the forward dynamics of well known, nonlinear, low dimensional dynamical systems. We further extend the results more practical, higher dimensional, robotic environments implemented in MuloCo. We use the D4RL dataset by Fu et al. (2020) to train Koopman autoencoder. Lastly we employ our proposed approach open-loop controller for locomotion tasks D4RL.
4.1 DYNAMICAL SYSTEMS
We use establsttbl dynamical systems as benchmarks for forward dynamics modeling. Despite their low dimensionality. these systems display interesting nonlinear dynamics, including multiple fixed points. Nonetheless, their low dimensionality enables the generation of informative visual representations in 2D/3D phase plots. We briefly review the environments used in this section.
Parabolic Attractor, adopted from et al. (2014); Brunton et al. (2016), a dynamical system with single fixed point at the origin, known for its closed-form Koopman embedding solution. Governed by the following equations:
(10) the system admits a solution that is asymptotically attracted to the parabolic manifold, given by T2 = x for μ < O. The Koopman embedding. 2, that adheres to globally linear dynamics, can be coded by augmenting the state with the additional nonlinear measurement of 23 = 27:
2 pmz. i2 A(z2 - z
0
21
[21
8 =
22
=
0
A
22
for
22
=
F2
(11)
23
0
2p.
L23
[2g
We set 1.0 and μ 0.1 and we sample initial conditions uniformly from 21.22 [-1, 1,,1]
Duffing Oscillator follows a nonlinear second order differential equation ă = - x² which represents model for the motion of damped and force-driven particle. This particular instance admits two center points at (z, . (+1, 0), and an unstable fixed point at the origin, (x, .) (0. 0).
Initial conditions are sampled uniformly from 21 [-2, 2] and Ø2 E [-1,1].
Lotka-Volte Volterra represents the population evolution of biological systems, based on predator- prey interactions, by the following equations:
è = 3a
d 5x 7222-
(12)
The system is known for its abrupt switch population growth and admits two fixed points, one at the origin (extinction), and center point at (æ1, x22 (7/5, a/B). We set 3 7 0.2 and uniformly sample initial conditions from æ1, 32 E (0.02.3.0.
Pendulum represents freely swinging pole. The initial conditions indicate the states from which the pole is released, deviating slightly from the inverted position by +10". The state consists of the angle and the angular velocity and report emror in radians.
5

---

Lorenz System is chaotic dynamical system (Lorenz, 1963). It features equilibrium points, some stable and some unstable, and is renowned for the "butterfly effect" arising from its sensitivity to initial conditions. The governing equations
follows:
o ( z1), 22 = 2p x3) 22.
21 Brz
(13)
We use the original parameters from Lorenz' 63 system. Initial conditions are generated by perturbing the point (0, I, 1.05) with Gaussian distrian- noise having standard deviation of 1.
To demonstrate data-efficiency, we applied our proposed Koopman Autoencoder to modest datasets of trajectories gathered from each of the aforementioned dynamical systems, comprising 100 trajectories for Lorenz systems and 50 trajectories for non-chaotic ones. For all dynamical systems except Lorenz' 63, the model is trained using the first 500 steps of the trajectories. Nevertheless, during inference, we unroll the models for up to 1000 steps, demonstrating the capability of our approach to accurately capture the underlying dynamics and generalize to unseen regions of the state space. We employed timesteps of 0.01 for forward integration in all environments, with the exception of Lorenz, for which we used a timestep of 0.02. We utilize embedding size of 128 for the dynamical systems.
Results are presented in Table 1. We conducted experiments involving both linear and nonlinear decoders, linear and nonlinear latent dynamics, as well as experiments with and without periodic reencoding. In the nonlinear latent dynamics setting we use an MIP, instead of K, to drive the latent state forward. We maintain consistent encoder and decoder capacities across different models applied to the same environment. We train a standard MLP for single step dynamics prediction as baseline, with capacity roughly equivalent to that of the encoders.
(a) Without reencoding.
Periodic reencoding consistently improves the accuracy the predictions. These improvements also extend to the setting where we permit nonlinear dynamics in the latent space, even though our method was motivated by learning linear latent dynamics. The only exception is the Parabolic Attractor environment (see Figure 3a). This is expected because this system admits simple feature transformation that achieves globally linear dynamics. The representation can be decoded, with a linear decoder (Equation 11) by simply selecting the first two elements, without the need for course correction. Furthermore, it should be noted that the best prediction results use linear decoder, rather than nonlinear one. This observation holds true, regardless of the type of dynamics assumed in the latent space or the specific environment. The results also demonstrate that Koopman autoencoders with linear dynamics and decoders, unrolled using periodic reencoding, consistently outperform widely used nonlinear dynamics models.
(b) With reencoding.
Figure 3: Parabolic Attractor.
4.2 D4RL: STATE PREDICTION
In this section, we train our proposed Koopman Autoencoder on continuous control tasks from D4RL benchmark by Fu al. (2020). Again, the goal is to predict the future states over long time horizons. We choose a number of locomotion tasks as our main testbed, namely the Hopper 12,, Hal HalfCheetahh-ee
environments. The curated datasets are generated by driving the simulated robots forward from a stationary position, using policies with varying degrees of optimality, i.e. expert, medi em-eep medi um, and medi replay. Each individual dataset comprises million transitions, organized as trajectories with maximum length of 1,000 steps. We use 80% of the trajectories for training and evaluate on the remainder 20% consisting of trajectories of length 300. Table 2 provides breakdown of the MSE over 300-step horizon.
and Walker2d-u
The D4RL benchmark differs from the dynamical systems discussed in the previous section in three fundamental ways: the dynamics are notably more challenging due the presence of collisions, the dimensionality is higher, and the systems are subject to control inputs. Here we use the extended version of the Koopman autoencoder model designed for controlled systems (see Section 3.... We use embeddings sizes of 512 and 256 for the state and action encodings, respectively. For training, we solely rely on the observations and avoid incorporating additional training signals, such as rewards.

---

MODEL
KOOPMAN (LINEAR LATENT DYNAMICS)
NONLINEAR LATENT DYNAMICS
MLP
DECODER TYPE
LINEAR
NONLINEAR
LINEAR
NONLINEAR
PERIODIC REENC. (X.
ENVIRONMENT
MSE over 100 steps
Parabolic Attractor
0.0205
0.0292
0.1465
0.0758
0.0739
0.0496
0.0727
0.0547
0,2674
Pendulum
0.0512
0.01M2
0.0648
0.0181
0.0288
0.0025
0.0242
0.0034
0.7442
Dufling Oscillator
0.1152
0.0112
0.1512
0.0512
0,0450
0.0012
0.0450
0.0112
0,4050
Lotka- Lotka-Volerrr
0.0113 X
0.0072
0.0145
0.0098
0.0128
0.0040
0.0112
0.0060
1.4450
Lorenz'65
12.569
7.265
19.051
88.565
ENVIRONMENT
AISE over 1000 steps
Pendulum
9.2021
0.1818
14,5800
0.6612
10.7184
0.0841
10.5832
0.29664
55.281
Dutling Oscillator
5.7851
0.5725
9.2451
22.445
Lotka Volterra
1.6292
0.3961
1,6203
0.2812
0.7261 X
0.2888
0.7281 X
0.4324
83,205
Lorenz'66
X
78.980
64.838
$9.262
54.793
133.509
Table I: Mean Squared Error of state prediction for a number of dynamical systems. The entries in the table are scaled by factor of 100x, except for Lorenz system. Evaluation is done via sampling unseen initial points from the dynamical system. The cross marks and check marks in the PERIODIC REENC. rows indicate, respectively, absence and presence of periodic reencoding mechanism during inference time. When enahled, the errors are reported by searching over reencoding periods of (10, 25, 50, 100). The cross marks in place of table entries indicate exploded values. Underlined values denote best performance for Koopman-based models, where bold numbers represent best performance across all models. We exclude the Parabolic Attractor environment when reporting errors over the 1000-step horizon since trajectories would almost perfectly merge onto the parabolic manifold and reach the origin in less time.
Similar to the previous section, a multi-step MLP is trained baseline. To ensure fair comparison, we set the size of the MLP to be equal to the combined size of the Koopman encoder and decoder. In our implementation, this equivalent MLP is implemented by setting the reencoding period during training and deactivating the loss terms associated with the Koopman autoencoder, specifically the "alignment" and "reconstruction" losses, optimizing only the "prediction" loss as the objective (see Figure 1). Refer to Appendix C for details and specifics.
We utilize training sequences with length of 100. The results presented in Table demonstrate that through periodic reencoding, we can consistently unroll our model over extended horizons, surpassing the training length (in this case, over a 300-step horizon). Furthermore, the use of nonlinear decoder proves to be crucial for accurate prediction. In experience, training the MLP becomes unstable for horizons longer than 10 steps, and its performance is significantly inferior. However, fulfills its role baseline.
MODEL.
KOOPMAN AUTOENCODER
I
MI.P
DECODER TYPE
LINEAR
NONLINEAR
PERIODIC REENCODING (X.V)
x
X
Environment
Dataset
MSE 300 steps
Hopper
expert
0.250 +: 0.05
0.079 = 0.03
0.353 + 0.05
0.012 + 0.00
0.561
medium-expert
0.486 0.07
0.102 0.04
0.719 + 0.10
0.015 0.00
0.692
medium
0.624 +
0.102 0.03
0.889
+ 0.08
0.008 + 0.00
0.533
full-replay
1.052 1.112
0.370 0.11
2.254 1 0.14 X
0.158 1 0.01
0.815
medium-replay
1.190 0.22
0.776 - 0.18
0.296 + 0.07
0.936
expert
0.645 0.05
0.602 = 0.09
0.07 0.622
0.227 + 0.03
1.359
medium-expert
1.141 0.05
1.076 = 0.08
0.813 +.13
0.391 0.07
L.481
medium
1.362 0.004
1.456 - 0.08
1.816 0.16
0.809 1 0.09
1.861
full-replay
1.262
1.252 = 0.17
1.668 + 0.17
0.816 + 0.14
1.994
W2
expert
0.364 0.07
0.302 = 0.08
0.544 + 0.05
0.072 + 0.03
0.285
medium-expert
0.755 1 0.02
0.602 1 0.08
0.796
1 0.10
0.198 1 0.08
1.295
medium
0.825 - 0.08
0.718 - 0.11
1.822 0.18 X
0.404 + 0.13
0.821
full-replay
1.676 - 0.19 X
1.413 = 0.13
0.867 1 0.15
1.291
redium-replay
2.379 0.20
X
2.077 1 0.26 1.917
Table 2: Mean Squared Error of state prediction for D4RL robotic locomotion tasks. Evaluation is done under a held-out set of trajectories. The cross marks in place of table entries indicate exploded or almost exploded values (large errors).
7

---

14-
(a) No REENCODING
(b) REENCODE EVERY STEP
(c) PERIODIC REENCODING
Figure 4: Phase portraits of the trained Koopman autoencoders (w/linear decoder) for unrolls of 1000 steps. We use different different reencoding schemes for unrolling the same model, namely (a) w/o reencoding, (b) reencoding every step. and (c) periodic reencoding (our proposed approach). The grey lines in the background represent ground truth phase portraits.
Figure 5: Mean Squared Error (MSE) over the unrolling horizon, under varying reencoding schemes. We use 100 freshly sampled (unseen) initial points for evaluation. 'reencode 0 legend indicates unrolling without ever reencoding. We observe that Periodic Reencoding robustly improves the quality of long-range modeling dynamics.
4.3 D4RL: SEMI-OPEN-LOOP CONTROL
In this section, utilize proposed model open-loop controller for locomotion tasks within the D4RL framework to showcase long-term stability, generalization capability, and prediction quality. This approach allows us to present another informative metric, the total reward achieved by a semi-open-loop controller (with sparse feedback). In this context, the model is trained using transitions generated optimally trained policy, specifically the expe rt datasets.
8

---

Mean Squared Error (MSE) for Hopper" environment over the unrolling horizon of 300 steps. under varying reencoding schemes. Periodic Reencoding generates more accurate trajectories compared to no reencoding and every-step reencoding schemes. Error plots for other environments be found in Appendix F.
Figure 6:
The objective to produce the sequence of states and actions that follow the initial state which provided input the model. We train a Koopman autoencoder, jointly with an independent decoder head trained using a behaviour cloning loss to output actions, which receives the encoded states input. We assess the quality of the predictions by evaluating the optimality of the replicated "expert" behavior. This is done by executing the generated sequence of actions in the environments and recording the total reward achieved. To underscore the stability of our approach, we ensure thar the model pluns the motion of the robot for the nexl 100 steps, by generating sequence of actions before receiving any sort of feedback form the environment. By removing feedback test time, i.e. by providing the ground truth state as input to the model only every 100 steps, we demonstrate the model's capability to generate stable motion plans, particularly when periodic reencoding is utilized.
MODEL
KOOPMAN AUTOENCODER
MLP
STATE FEEDBACK (X, V)
x
X
X
V
PERIODIC REENCODING (X, V)
Environment
Dataset
total reward until termination
Hopper
expert
18.40 6.2
53.5 + 14.5
18.9 5.9
96.05 + 4.3
HalfCheetah
expert
15.06 + 5.3
64.2 + 12.9
19.5 + 7.9
82.91 + 5.9
2a
expert
18.46 + 8.1
61.9 14.2
2..4
98.73 = 1.5
Table Semi-open-looy reward achieved. The total rewards are normalized according to D4RL random and expert scores. The Koopman Autoencoder uses nonlinear decoder.
Table 3 provides the results for the designed open-loop control problem. As baseline, we train an MLP that takes the state as input and outputs the next state along with the optimal action. The MLP can then be unrolled autoregressively test time similarly an open-loop fashion. Furthermore, for comparison, we train a standard behavior cloning (BC) model that we run without state obfuscation, and allow to observe state every step. This the upper bound of performance of the policy run with state obfuscation.
We demonstrate that our approach capable of making sufficiently accurate predictions well into the future when the state is withheld from the model over 100-step horizons. This is evident from the agent's ability to sustain its motion without falling, relying solely on preplanned motion.
CONCLUSION
Our study of Koopman autoencoders for modeling nonlinear dynamical systems lead us to explore various inference schemes for generating predictions over long horizons. We showed that generating trajectories exclusively in the latent space presents two potential difficulties: (i) the inability to capture switching behaviour between multiple fixed points, and (ii) violation of the existence and uniqueness theorem of initial value problems. We showed, through theory and experiment, that trajectories generated with reencoding do not suffer from these limitations. Finally introduced periodic reencoding as method that bridges the gap between no reencoding and reencoding at every step, and achieves the best results in practice.

---

ACKNOWLEDGMENTS
We would like to thank Daniel Worrall for careful review of the manuscript and useful feedback. We would also thank Peter Battaglia and Hugo Larochelle for their prompt approval the work for publication.
REFERENCES
Martín Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. CoRR, 00s//511.064466 2015. URL http:///axixx org/ hass//ab///// 1511. 06464.
Omri Azencot, N. Benjamin Erichson, Vanessa Lin, and Michael W. Mahoney. Forecasting sequential data using consistent koopman autoencoders. In Proceedings of the 37th International Conference on Machine Learning. pp. 475-485, PMLR, 2020.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances neural information processing systems, 28, 2015.
James Bradbury. Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miin and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http://githrn com commggogle/////ooooooooooo jax.
Steven Brunton, Bingni W. Brunton, Joshua Proctor, and J. Nathan Kutz. Koopman invariant subspaces and finite linear representations of nonlinear dynamical systems for control. PLOS ONE, 11(2)e0150171, feb 2016. doi: 10. Io.mmall. 0150171. URL https: org/10. 337182Fjournal pone, 0150171.
Steven L. Brunton, Marko Budišić, Eurika Kaiser, and Nathan Kutz. Modern koopman theory for dynamical systems. SIAM Review, 64(2):229-340, 2022.
Anthony Frion, Lucas Drumetz, Mauro Dalla Mura, Guillaume Tochon, and Abdeldjalil Aissa El Bey. Leveraging neural koopman operators to learn continuous representations of dynamical systems from scarce data, 2023. arXiv preprint.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020. arXiv e-prints, April 2020.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state-space layers. In NeurlPS: Proceedings of the 34th Neural Information Processing Systems Conference, December 2021.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022. URL https:
poster/6959.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. arXiv preprint 151.:1811.045511 2018.
R.E. Kalman. On the general theory of control systems. In Proceedings of the First International Congress on Automatic Control, pp. 481-492, 1960.
B. O. Koopman and J. Neumann. Dynamical systems of continuous spectra. Proceedings of the National Academy of Sciences, 18(3):255-263, 1932. doi: 10. 1..107/pas..1...222 URL https: www. pnas. org/ /org a 1tt./111..17//aa.1.1 18. 255..
Bernard O Koopman.
Hamiltonian systems and transformation in hilbert space. Proceedings of the National Academy of Sciences, 1755):315-318, 1931.
10

---

Martin Wilhelm Kutta. Beitrag zur näherungsweisen integration totaler differentalgleeihungen. Zeitschrift für Mathematik und Physik, 46(1):435 453, 1901.
Yueheng Lan and Igor Mezić. Linearization in the large of nonlinear systems and koopman operator spectrum. Physica D: Nonlinear Phenomena, 242( (1:223) 2013.
Edward N. Lorenz. Deterministic Nonperiodic Flow. Journal of Atmospheric Sciences, 20(2): 130-148, March 1963. doi: 101111/1122000000000000000000000
Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam. CoRR, abs/ 171105001,, 2017. URL http:///eeeee //ttpi///xxvv.../// org/ 11w.. 05101.
Bethany Lusch, J. Nathan Kutz, and Steven I.. Brunton. Deep learning for universal linear em- beddings of nonlinear dynamics. Nature Communications, 9(1), nov 2018a. doi: 10.1038/ s41467-018-07210-00 URL https:/ 0001 41187-018 111---11--1-0001------- 07210-0.
Bethany Lusch, Nathan Kutz, and Steven L Brunton. Deep learning for universal linear embeddings of nonlinear dynamies. Nature communications, 911:4950, 2018b.
Xubo Lyu, Hanyang Hu, Seth Siriya, Ye Pu, and Mo Chen. Task-oriented koopman-basee control with contrastive encoder. In 7th Annual Conference on Robot Learning, 2023.
Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL https: //openreview...e.e//m./e.....eee net wwr.moiduuidiooeoo Hidnniddd--Wwlcc
Arnab Kumar Mondal, Siha Smarak Panigrahi, Sai Rajeswar, Kaleem Siddiqi, and Siamak Ravan- bakhsh. Efficient dynamics modeling in interactive environments with koopman theory, 2023.
Zahra Monfared, Jonas M. Mikhaeil, and Daniel Durstewitz. On the difficulty of learning chaotic dynamics with mns. In Advances in Neural Information Processing Systems, vol- ume 35, 2022, URL https: /ttpss://ppppeeee /wtppraapessse//////eeeeeeeeeeee nips wttppccppppppeeeeeeeeeeeeeeeeeee httpppappp/////// 20/22/22222/////// hash/ 495e55
html.
Jeremy Morton, Antony Jameson, Mykel J Kochenderfer, and Freddie Witherden. Deep dynamical modeling and control of unsteady fluid flows. Advances in Neural Information Processing Systems, 31, 2018.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 26670-26698. PMLR, 2023.
Barak A. Pearlmutter. Dynamic recurrent neural networks. Technical Report CMU-CS-90-196. Camegie Mellon University, School of Computer Science, December 1990. URL https: mural www.ootthnivvess......... 1e/50/5////////
Carl Runge. Über die numerische auflösung von differeenialgleechunge.. Mathematische Annalen, 46(2):167-178, 1895.
Christof Schütte, Péter Koltai, and Stefan Klus. On the numerical approximation of the perron- frobenius and koopman operator. Journal of Computational Dynamics, 3(1):1-12, sep 2016. doi: N..3934/jcd.201600 URL https: //doi. otreiioorg1100 191.111..33333331 201 6003.
Jimmy T.H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling. International Conference on Learning Representations, 2023. ICLR 2023, Top 5
Adam Trischler and Gabriele MT D' Eleuterio. Synthesis of recurrent neural networks for dynamical system simulation. Neural Networks, 80:67-78, 2016. doi: 10.10166 10.10116/ w/tp:///eneet00000 www..eeee.211.4.000....

---

Jonathan H. Tu,, Clarence W. Rowley, Dirk M. Luchtenburg, Steven L. Brunton, and J. Nathan Kutz and. On dynamic mode decomposition: Theory and applications. Journal of Computational Dynamics, 1(2):391 2014. doi: (0.3994/jed.20144..399. URL https: 399382FFJe .tt.:icdd2114..0...
201 0.0014.25 1.911.
A. Tustin. A method of analysing the behaviour of linear systems in terms time series. Journal of the Institution of Electrical Engineers Part IIA: Automatic Regulators and Servo Mechanisms,, 94:
130-142(12), May 1947. ISSN 2050-5523. URL https: / org/ http//onteet//ooraallloo
theiet.
10100099/1100 1947.00220 0020.
Pantelis R. Vlachas, Georgios Arampatzis, Caroline Uhler, and Petros Koumoutsakos, Multiscale simulations of complex systems by learning their effective dynamics. Nature Machine Intelligence, 2022.
Matthew O. Williams, Clarence W. Rowley, and Ioannis G. Kevrekidis. A kernel-based approach to data-driven koopman spectral analysis, 2014. arXiv e-prints.
Matthew O. Williams, Ioannis G. Kevrekidis, and Clarence W. Rowley. dlata-driven approximation of the koopman operator: Extending dynamic mode decomposition. Journal of Nonlinear Science, 25(6): jun 2015. doi: 10. 00010077600332--1---555555. URL bittp:://di..rre
12

---

A IMPLICATIONS FOR ALL LATENT MODELS OF DYNAMICAL SYSTEMS
Equations and can be generalized beyond linear latent dynamics to any deterministic latent space dynamics. Simply let K be function of instead of constant matrix giving rise to latent space dynamics of å = K (z), with solution z(t) hit, o(zo)). Without reencoding the solution in state space becomes:
z +b (h(t, ((x))))
(14)
Note that because is the solution to a dynamical system in 2, its phase lines would never cross in latent space. However, there is no such guarantee in space for arbitrary vb. Reecnoding gives rise to the following dynamical system: a == (=J)(o(z))K((or))
(15)
where K is now simply function, taking = - ofx) as input. Naturally, the solution would satisfy the uniqueness theorem of initial value problems, i.e. phase lines would not cross. These findings empirically validated generalize the nonlinear case by learning a nonlinear latent dynamics (parameterized by an MLP), see the 'NonLinear' Latent Dynamics' column in Table 1.
B KOOPMAN THEORY OVERVIEW
Let us consider an autonomous dynamical system defined over an open set D of TR" by system of first order differential equations:
x = f(x)
(16)
This system induces a flow & D x R R + D which, for a given time f, an observed value of the state of the system and duration At, maps to the state of the system evolved by said duration
At. f(x,) dr Xt+
(x)+ =
(17)
One of the challenges in modern dynamical system theory revolves around establishing transfor- mations which maps such nonlinear dynamics into linear ones, as there exist myriad of tools and frameworks facilitating the study of linear systems, in contrast to their nonlinear counterparts.
In his seminal work Koopman (1931) proposed an alternative perspective of general nonlincar dynamics in which he considered measurement functions & D - R s.t. & E L the set of square integrable functions w.r.t. Lebesgue's 's measure which constitutes Hilbert space whose metric determined by the inner product (1p, 1/) Jo (x)))))())))(()) A measurement function maps the actual state of the system to the measurement or observable data to which one would have access. He introduced an operator K. acting upon the set of measurement functions and which advances the measurement such that
V 2 VX D, Vt, IR s.t. t> 8, Kx-)(1)(x) + (x)-,(())
(18)
is linear by definition and the family of operator {Ko-sehee constitute one-parameter group in Hilbert space, which admits infinitesimal generator the Lie derivative operator T such that
Ne-ee (x) - (x((t)
(())(x(t)) lim 7-00
(19)
lim (x((t + T)) (p(((t)) di ((x((0)) (V)) f(x(t))
The resulting linear differential equation on measurement functions is of little use in practical settings, however, it is of infinite dimension. To cireumvent this, it is customary to resort to orthogonal decomposition of the Hilbert space of measurement functions X2 into subspaces invariant by application of the Lie (and thus Koopman) operator (Brunton et al., 2016; Morton et al., 2018).
2
k
(20)
A=
13

---

By truncation, it possible to operate invariant subspace 5 θ Cx. thus providing finite dimensional representation of the operator ' (resp. X)), restricted to said subspace, into a matrix T (resp. K) acting on the vector space RR. The restriction of measurement function 6 thus fulfils the linear dynamics
d di dt
[ -TP"(x(1)
(21)
or equivalently
-*(x(t))=<pp(TT = (t to)) (²(x(tt)) = ((t--o) (x((t))
(22)
C IMPLEMENTATION DETAILS
C.I. STATE PREDICTION TASKS
To minimize the alignment loss, the encoder may generate latent codes with arbitrarily small values. We regard these solutions as degenerate, and to discourage this behavior, we normalize the weight columns of the decoder. Moreover, we employ AdamW optimizer by Loshchilov Hutter (2017) with a slight weight decay of value 1e-4 and learning rate of 1e-4. We utilize a custom learning rate of 1e-5 for the dynamies components, which are either the Koopman matrices or 2-layer MLPs. This encourages the encoder to follow the dynamics prescribed by Koopman and stabilizes training. We use the continuous Koopman parameterization and apply bilinear descretization (Tustin, 1947) whenever control inputs are present. In the absence of control inputs, allow the model to unroll by solving an initial value problem (IVP), using more sophisticated integrators implemented by JAX (Bradbury et al., 2018). During both training and inference, cemperimentain immentall.......mm...........e.l ode. odeint () which relies on an adaptive stepsize (Dormand-Prinnn Princc) Runge-Kutta integrator (Runge, 1895; Kutta, 1901). We find incorporating the prediction loss to be detrimental to the overall performance when training on the dynamical systems listed in Section 4.1. Inspired by the model described Appendix D. we use small sparsity inducing L1 loss, le-3, applied to the Koopman embeddings to encourage region-dedicated dynamics, and use ReLU activations for all experiments to ensure sparse activations. The encoders used for dynamical systems and D4RL are 4-layer and 6-1ayer standard MLPs, respectively. We use training sequence lengths of 10 and 100, for the dynamical systems and D4RL state prediction tasks, respectively. Moreover, observe that using reencoding training schemes of 20 and 50 steps are beneficial for most of D4RL state prediction tasks, compared to training without reencoding. Interestingly, presence of control inputs in D4RL. dataset makes the training more stable over longer horizons, compared to the dynamical systems we use as henchmark. This is because in the absence of control inputs, the latent state can be arbitrarily scaled up by the eigenvalues of the K matrix. We find that normalizing the observations and actions as a preprocessing step negatively impacts performance, especially for D4RL tasks. We choose an embedding size of 128 for all dynamical systems, 512 for D4RL states, and 128 for D4RL action embeddings. The stepsize is trained in log space, initialized from that of the environment.
make heavy use of the
C.2 D4RL: SEMI-OPEN- CONTROL
In this subsection review the implementation details that pertain specifically to the open-loop control task. Given that the actions in the expert datasets are generated via a fixed expert policy, which is function of the states, denoted W, 7 T(xx). we can transform the dynamics from their original form, x++1 = F( (.), more compactly into, Iz+1 F( π (x²)) = F, (rr,). Therefore, we employ the standard Koopman formulation, i.c. 2t+1 K2, which does not involve exogenous controls (see Section 3.1). Due to the absence of controls as direct inputs to the model, dictated by the nature of this problem, the training process becomes unstable for long sequences, in comparison the DARL state prediction tasks. We train separate decoder head for action prediction, which takes the encoded states as input. We utilize training sequences with a length of 20 steps, and likewise, employ periodic reencoding scheme of 20 steps during test time. After the first 100 steps taken by the agent in the environment, we restrict ourselves to observing the ground truth state values only at intervals of 100 steps, hence the designation Semi-Open-Lae we utilize precise adaptive-steptl stepsize integrators implemented in JAX, i.c. odeint (), both during the training and inference stages.
Additionally, for improved prediction accuracy,
14

---

D SWITCHING DYNAMICS
To explore another limitation of trajectories generated without reencoding, it is instructive to study the special case where dim((z) dim(æ) and vxx)) Wz, i.e. the decoder an over-complete dictionary, e.g. a linear operator. Therefore the trajectories generated without reencoding correspond to linear projection of the solution of a linear dynamical system, e.g. xxt) We 6(x)). Such trajectories exhibit limited expressiveness; for instance, they cannot capture transitions between multiple distinct fixed points. For example, consider scenario where x E IR. governed by NLDS with multiple fixed points, such as the Duffing Oscillator. Let z E TR². It is possible for to represent multiple, distinct, attractors of x but not the switching behaviour between them that occurs in the Duffing oscillator for some initial conditions, Let these attractors be located in two regions of state space, R and R2. Now, let z = [21, 22, 23, 24] and:
d(x) = {[ 22, 0.
if E R1, [0, 23, 21]
if R2)
(23)
These can be interpreted sparse codes with distinct supports corresponding to R1 and R2, inferred using øb Furthermore, if we let be block diagonal then:
K
0
Kg
2.
(24)
Phase space trajectories are synthesized by linearly combining the trajectories the latent space.
11
-2
W2
if Ri:
Z3
W3
-
24
W4
if E R2
(25)
The dynamics of x can be approximated by two distinct linear systems. Furthermore, the columns of W can shift the origin to represent two distinet fixed points. Note that, without reencoding (Equation 8), the dynamics of initial condition 2g = ((x)). for all time. In other words, this scheme does not allow for switching between linear systems, represented by Ki and K2, after t in Figure 4 (a), the dynamics captured without reencoding resemble the phase diagram of the superposition of several LDS. In contrast, trajectories generated with reencoding, have the capacity to switch between supports (non- elements of 2) and therefore switch between the dynamics defined by K₁ and K2. Indeed was shown that non-linear systems can only be linearized within basin of attraction of fixed point (Lan & Mezić, 2013), implying that the best we can hope to achieve is to partition the phase space into regions each which can be well approximated by LDS, as illustrated in the example above.
restricted to follow the same linear system, determine by the
Indeed this is what observed
E EFFICIENCY
Producing full trajectories for T discrete points in nonlinear dynamical system of the form T++1 F(x,) requires the repeated recurrent application of the nonlinear transformation F. For the Koopman Autoencoder, trajectories for T discrete points are first generated in the Koopman linear space ž4+1 K2 + Lv and followed by a time-agnostic decoding step = (zz). The linear recurrence in the latent space allows for highly parallelizable unrolling of the predicted sequence using parallel scans (Martin & Cundy, 2018). Compared to nonlinear recurrences, training time is greatly improved. Periodic reencoding does add a nonlinear step to the Koopman autoencoder recurrence. However, since we decode and reencode periodically every k steps, we can parallelize the predicted trajectories for i discrete points and still obtain much improved training efficiencies. Furthermore, considering that our method can provide accurate predictions over short horizons (where reencoding not needed), valuable in situations requiring fast look-aheas capabilities, such as model. predictive control (MPC) or ne- step value function bootstrapping in Reeinforcement Learning.
15

---

F ADDITIONAL RESULTS
environements Hopper-ys r- error plots can be found in the main text. For all D4RL, experiments, during training time, we either do plots correspond to the best performing trained models reported in Table 2. Periodic Reencoding consistently achieves the best results and is robust to the reencoding period.
Figure 7 show the eπror plots for Hal f2 and Walker2d-v2
reencoding all, we use reencoding periods of 20 50 steps. The
Figure 7: Mean Squared Error (MSE) for HalfCheetah" and Walker2d environments over the unrolling horizon of 300 steps, under varying reencoding schemes. The results are for Koopman Autoencoder models.
Figure 8 shows the error plots of the Koopman Autoencoder models that DO NOT employ reencoding (periodic reencoding) during training time. We still observe drift latent space despite the absence of periodic reencoding training time. Employing periodic reencoding test time still mitigate the drift and generate stable, long unrolls.
Figure 8: Mean Squared Error (MSE) for D4RL environments over unrolling horizon 300 steps, under varying reencoding schemes. The plots helong to models that DO NOT make use reencnding during training. The errors reported in Table 2, under PFriOdIC REENCODING: (X), corresponds to the blue lines.
16

---

Experiments with Structured Koopman Matrix. Additionally, we attempted experiments with a diagonal Koopman transition matrix but were unsuccessful. We suspect that this lack of success may be due to additional requirements imposed by the Koopman linear space. Achieving complete disentanglement of features in this context might necessitate an exponentially larger latent state dimension. In effort to attain training stability and mitigate the issue of scaling up and down the outputs in linear dynamics steps, we conducted experiments using a skew-symmetric matrix, denoted as K. Skew-symmetric transition matrices perform rotations without scaling the space. The experiments with the kkew-symmetri matrix yielded unsuccessful results, similar to those with the diagonal matrix. We speculate that this lack of success may be attributed to a lack of expressivity.
Low Dimensional Latents and Deep Latent-to-Latedt Mappings. So far, we have assumed high-dimensional latent space, where dim((z) >> dim(z), as in theory, Koopman operator is an operator. This enables to learn rich representations of the dynamics, with the information bottleneck being the linearity requirement by Koopman. We conducted experiments using lower dimensional state embeddings, dim(x)) > dim(=), and deep latent-to-latent transition networks. In this context, the information bottleneck more closely resembles that found in tranditional, "static" autoencoders, where the projection into the latent space represents form of compression. We observed extreme instability when unrolling over long horizons, resulting in poor prediction quality, We also discovered that Periodic Reencoding fails to improve the stability of these models.
Single Step Koopman Objective. Rather than training the Koopman objective over multiple steps, we focused on minimizing the losses as defined in Equation 7 at individual timesteps. We implemented this approach with the aim of perfectly aligning the training objective with the inference time scenario, where reencoding occurs at every step. However, this resulted in high inference- -eree instability and subpar performance, even worse than the MLP baseline, regardless of the utilized inference mechanism, e.g. with or without reencoding. This suggests that the Koopman objective is indeed leveraging its extensibility within basins of attraction, as indicated by Lan & Mezić (2013).
G CONNECTION TO RNNS
There exists an extensive body of literature on training RNNs on nonlinear dynamical systems (Pearlmutter, 1990; Trischler & D' D'Eleuterio, 2016; Vlachas et al., 2022). A prominent application of these models is in Model-Based Reinforcement Learning, where they function as "World Models." For instance, Hafner et al. (2018) introduced Recurrent State Space Models (RSSM) trained on image inputs. These models are often used for short-term future predictions while learning value function. In contrast, our work facilitates long-range predictions into the future. Koopman autoencoders can be regarded as RNN models with a single linear recurrent layer. From an efficiency perspective, linear recurrent units are more ad advantageous as they can be executed in parallel rather than sequentially. Furthermore, from a representation learning perspective, one could argue that linearly evolving features extracted from NLDS would result in more meaningful and interpretable representation of the system in question. State-space Models (SSMs) (Gu et al., 2022; 2021) can also be regarded as linear RNNs, achieving impressive performance in modeling long-range dependencies. recent finding by Orvieto et al. (2023) highlights that linear RNN layers surpass tuned nonlinear RNN variants in performance. This result is also partially attributable the fact that the behavior of linear layers can be designed and engineered more effectively, thanks to the well-studied nature of this problem. This enables effective regulation of linear layers, ensuring the stability of the training process over long sequences. A similar approach was previously adopted by Arjovsky et al. (2015), wherein the eigenvalues of the RNN transition matrix were constrained to lie the unit circle. The drawback of imposing structure and constraints on the transition matrices is that it results in loss of expressivity (Monfared et al., 2022).
17


Reviews:
[{"review_id": "eoqJP45y0f", "reviewer": "Reviewer_QVMC", "summary": "The authors of the study investigated the use of Koopman autoencoders to model complex, nonlinear dynamical systems. They explored different approaches for making long-term predictions and showed that generating predictions solely within the latent space had two potential challenges: 1) It was unable to effectively capture the behavior of multi-stable systems that switch between multiple fixed points. 2) It violated the existence and uniqueness theorem of initial value problems, which establishes the uniqueness of solutions in dynamical systems. To address these challenges, the authors introduced a technique called \"periodic reencoding.\" This method strikes a balance between not using reencoding at all and reencoding at every step, and it was found to yield the best practical results in their experiments.", "questions": "Please see above!", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written, featuring a fascinating main idea, and employing innovative and effective methods with good results.\n\n- The overview provided in the appendix is very helpful and is written in a precise and technical mathematical style.\n\n- The visualization provided in Fig. 1 is very helpful in providing a clear illustration of their method.", "weaknesses": "**A minor point**:  Regarding the sentence \"This characteristic becomes especially significant when the dimensionality of the latent space is substantially larger than that of the original space, i.e., n >> d, ...\", why is this case considered more important? While it's true that a mapping to a larger dimensional space is not surjective, as discussed in the paper, it seems that the lack of injectivity of the map is more critical. So, since a mapping to a smaller dimensional space is not injective, wouldn't the case of n << d be more problematic?\n\n**A few questions**:\n\n1. When it comes to learning chaotic dynamics, what are the advancements offered by the proposed method compared to using RNNs with teacher forcing, as presented in the following papers?  \n\n[1] F. Hess, Z. Monfared, M. Brenner, D. Durstewitz, Generalized Teacher Forcing for Learning Chaotic \nDynamics, Proceedings of Machine Learning Research (ICML 2023), (2023). \n\n[2] J. Mikhaeil, Z. Monfared, D. Durstewitz, On the difficulty of learning chaotic dynamics with RNNs, Advances in Neural Information Processing Systems (NeurIPS), (2022).\n\n2. Regarding eq. (9), it apears that the Jacobian $J_{\\psi}$  of  $\\psi$ is assumed to be a constant matrix. If this is the case, how can one ensure that it remains a fixed matrix? It may vary depending on the specific points at which it is calculated and could differ across different points.\n\n3. How effectively would the proposed method perform in capturing the dynamics of a multi-stable system with a stable fixed point and a periodic attractor, or possibly with two periodic attractors?\n\n4. How changes in the hyperparameters $\\Delta t$ or $k$ can impact the results? Does an optimal value always exist for that?\n\n5. Why do the authors use an MLP instead of K (e.g., on page 6)? How does the choice of activation function affect the performance of the MLP? Were other activation functions considered, and if so, how did they compare to the chosen function? Can the MLP be replaced with other types of neural networks, such as convolutional neural networks or recurrent neural networks? How would this affect the performance of the proposed framework?\n\nI am happy to increase my score if the authors could address my concerns.", "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": "DxB7yvnTrO", "reviewer": "Reviewer_FYpR", "summary": "This work proposes an autoencoder model for dynamical system data based on Koopman theory; observed variables are mapped to the \"Koopman feature\" space in which the state evolution is linear.  The authors proposed a periodic re-encoding scheme that avoid the issues of trajectory crossing in observation space and the inability to model switching dynamics.  The authors demonstrated that the proposed method with re-encoding leads to improved state prediction and controlling.", "questions": "See above.\n\n## Post-rebuttal update\n\nThank you for your response.  I am raising the score as my concern is largely addressed.  I remain uncertain about the interpretation of the RL experiments; while they have indeed demonstrated the utility of the proposed method in the Koopman framework, the lack of baseline comparisons makes it somewhat difficult to put them in context, i.e. to what extent it is addressing an issue in an application where the Koopman framework is of interest.  As I am not very familiar with this area, I will defer the judgement of this issue to the other reviewers.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper is mostly well-written.\n- The proposed method demonstrates promising empirical performance.", "weaknesses": "- I feel the theoretical benefits of re-encoding are not well explained.  It appears that re-encoding should be most useful in eliminating undesirable local optimas; if the model is correctly specified, the AE objective should have the same global optima with or without re-encoding.  It is thus not clear whether it is a principled solution to the switching dynamics issue, which is one of the two motivations for this work; if the dynamics cannot be (globally) linearized the model is misspecified in either case.\n\n- While the experiments demonstrate the benefits of re-encoding for autoencoding approaches, it is not clear why we should restrict to such approaches for control. There is also a general lack of comparison with baselines, e.g. with Mondal et al (2023).\n\n- As a minor point, the reference to deep state-space models such as S4 is slightly misleading as they are often applicable to stochastic systems, whereas this work studies deterministic systems.", "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": "g5fTypDiRp", "reviewer": "Reviewer_Fg8g", "summary": "This paper introduces the concept of periodic reencoding into the field of Koopman autoencoders. The idea is that, rather than propagating the dynamical system solely in the latent space, we periodically re-encode the true-space state $x_t$ to the latent space. This overcomes possible issues with $\\phi$ and $\\psi$ not being exact inverses (due to lack of bijectivity).", "questions": "- has equation (2) been used in previous papers? I see that Lusch et al. (2018) used something similar. It would be helpful to precisely clarify.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- the concept of periodic re-encoding is clear and simple\n- it leads to across-the-board improvements in performance", "weaknesses": "- any underlying mechanism governing periodic re-encoding is not discussed. For example, are there deficiencies in the encoder or decoder that, if adjusted, remove the need for reencoding, or change the optimal period? \n- how is the optimum period chosen? Since every step reencoding is worse in terms of performance than periodic reencoding, this is not a simple tradeoff in terms off costs\n- what causes do we hypothesise for \"long term drift\"? Perhaps during training the encoder does not learn to \"use\" all the latent space. Thus, there exist points in latent space with a decoding that is not properly trained. This is pure conjecture, but it would be beneficial for this paper to investigate these issues.\n- text is not legible in some figures (Fig 3, 4, 5 font size too small, Figure 6 more extreme)", "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": "Xt5OliIXnT", "reviewer": "Reviewer_Woag", "summary": "This work we study Koopman autoencoders to model nonlinear dynamics, specifically for long term state forecasting. The authors propose Periodic Reencoding, an inference-time mechanism that incorporates far future (multiple-step ahead) loss for faithfully capturing long term dynamics. The proposed method is validated both analytically and empirically via experiments in low and high dimensional NLDS.", "questions": "It's worthy discussing the \"optimal\" $\\Delta t$ and its relationship to the dynamics. For example, Fig. 5 shows the best are the middle ones.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The proposed Periodic Reencoding improves long term forecasting of deep Koopman autoencoders. \n\nIt does not require long sequence training. \n\nThe manuscript is well written.", "weaknesses": "The Periodic Reencoding works well especially with oscillator dynamics. It is likely related to the fact that those dynamics is periodic. It would be interesting to show how it works with continuous attractors.", "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": "wRyE38WsmT", "reviewer": "Reviewer_f4vn", "summary": "Given a fully observable deterministic nonlinear dynamical system described by $dx/dt = f(x)$ such as Lotka-Voltera predator prey population models or chaotic pendulums, Koopman theory states there exists a mapping to a (typically much higher, possibly infinite, dimensional) space $z = g(x)$ where the dynamics are linear, fully described by $dz/dt = Kz$. (I am loosely reminded of the use of kernels in SVMs, a binary classification dataset where classes are not linearly separable can be mapped to a higher dimensional space where classes are linearly separable).\n\nThere is an active body of work using the Koopman framework where the mapping function $z = \\phi(x)$ is a learnable neural network trained with a corresponding inverse mapping $x = \\psi(z)$ which together form an encoder – decoder pair in an autoencoder architecture (albeit where latent $z$ dimension may be lower or higher than the data $x$ dimension), together with this autoencoder, the method learns the transition matrix across time in state latent space $z_{t+1} = K z_t$.\n\nFrom an initial datapoint $x_0$, time series prediction is performed by encoding the initial point $z_0=\\phi(x_0)$, applying the latent state transitions, $z_t = K^tz_0$, and decoding the latent states $x_t = \\psi(z_t)$. The authors state that this inference procedure can fail for long term prediction as latent to trajectories that do decode to a sequence of  “drifting” datapoints. Prior works have proposed heuristic solutions such as training on long sequences and restricting the eigen values of the latent space transition matrix K. The authors propose an alternative heuristic they name “periodic re-encoding” where, at every kth transition also incorporates a decoder – encoder pass, i.e. \n\n$z_{t+1} =$\n\n$ K\\phi(\\psi(z_t))$ if $t \\text{mod} k ==0$ \n\n$Kz_t \\quad \\quad \\quad$ otherwise\n\nThe authors perform experiments on simple well known NLDSs with and without periodic reencoding and make a strong case for the efficacy of periodic re-encoding. The authors apply the method to more challenging datasets from D4RL library and again demonstrate the improvement in performance as a result.\n\n\n## Update After Rebuttal\nI sincerely apologize for missing the response period, I hope the authors can still read this!\n\n- my initial judgement on the value of Koopman representations seems misguided, and the authors and other reviewers responses have significantly increased my appreciation and I have raised my paper score to a weak accept whilst also reducing my confidence score appropriately.\n\n- I now understand the bilinear approach, I did not see the inverse matrix (it as a half step forward and a half step backwards taken in reverse)", "questions": "- For each problem, can you state the dimensions of $x$ and $z$ in the main paper. (I couldn’t see them in Appendix B)\n- given linear latent transitions are worse and periodic reencoding is more general than Koopman autoencoders, what is the justification for the authors to focus on Koopman representations? What is the benefit of the Koopman framework over any other autoencoder based time series model?\n- Equation 6: it appears the transition matrix is applying ½ of a forward time step and then ½ of a backward time step? I would understand if the transition matrix was made of many incremental time steps e.g. $\\textbf{K}=(I + \\frac{\\delta}{n}K)^n$ where $n =2$ (and consequently $\\textbf{K}=exp(\\delta K)$ as $n>\\infty$)\n- what is the benefit of periodic re-encoding, clamping (high dim) latent representations to the latent data manifold or adding noisy perturbations?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- the experiments appear to be an ablation study on the effect of periodic re-encoding and show a clear benefit\n- the inclusion of both toy and more complex examples is much appreciated.\n- there are many works combining autoencoders with latent space transition models. Presumably periodic re-encoding can be applied to any model with the same autoencoder + latent space transition architecture, is it more general than the authors present? \n\n- I felt it was generally well structured and written, although a few minor writing comments mentioned below.", "weaknesses": "## Update After Rebuttal\nSee Summary\n\n# Original Review\n\nFor context in the follow up discussion, I would like to state my intuition of Koopman theory is that any finite dimensional nonlinear system can be “lifted” into a higher dimensional space where dynamics are linear, in other words a complex low dimensional nonlinear transition function can be replaced with complex nonlinear high dimensional mapping and a simple transition function, we have just moved the “non-linearity-ness” from one place to another and Koopman theory states that this is universally possible for any NLDS, although in practice this may be very difficult or even impossible in finite dimensions.\n\n## Major Points\n- If I understand correctly, when $dim(z) > dim(x)$, the data points encode to a $dim(x)$ manifold embedded in the z-space. One interpretation is that the re-encoding takes latent points that have drifted away from this from the latent data manifold and puts them back on. However, if this were the case, re-encoding every step would have the best results.\n- I struggle to agree with the conclusions drawn by the example given in Appendix C.  I view Appendix C as a special case of failure to specify good mappings. The claim that periodic reencoding solves the problem I find strange. If I understand correctly, the perturbations introduced by the decoder-encoder pass allow the latent state to jump between basins of attraction of each fixed point? Presumably these perturbations are just noise and an artefact of the (arbitrary) precision achieved by model training convergence.\n- what is the benefit of periodic re-encoding, clamping latent representations to the data manifold or adding noisy perturbations?\n- the authors include non-linear transitions as a baseline, while it is nice to have more baselines, as I understand, linear transitions is the only reason to use Koopman theory. As expected, the (strictly higher capacity) non-linear transition outperform the liner transitions.\n- Given the re-encoding is not specific to Koopman, and that Koopman (linear) transitions predictably perform worse than non-linear transitions, this paper’s focus on the Koopman framework seems unnecessary.\n- there are many many autoencoder + latent space transition models that vary mainly by application specific architectural differences (e.g. pixels/graph/vectors require CNN/GCN/MLP encoder) and their training losses (L2, ELBO, Cross entropy etc), and a quick google for \"latent dynamics\" and \"deep kalman\" yields many results, many that have linear latent transitions\n  - https://papers.nips.cc/paper_files/paper/2017/hash/7b7a53e239400a13bd6be6c91c4f6c4e-Abstract.html\n  - http://www.approximateinference.org/2015/accepted/KrishnanEtAl2015.pdf\n  - https://arxiv.org/abs/1811.04551\n\n\n\n## Minor Points\n- Section 2, first paragraph does not help unfamiliar people understand Koopman theory or MDM (and familiar people do not need to be told). I would suggest dumbing this down as the level of Koopman theory required for this paper can be explained very easily.\n- Bibliography: many of the references have titles, names and dates but are missing journal/conference information\n- Appendix E is a nice addition, yet it feels somewhat arbitrary to cite a handful or RNN papers when the time series modelling literature is vast and the autoencoder + latent dynamics literature is I feel is more closely related this work than RNNs.", "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)