Disentangled generative models for robust dynamical system prediction
Authors: Stathi Fotiadis
Shunlong Hu
Mario Lino Valencia
Chris D Cantwell
Anil Anthony Bharath
Venue: ICLR 2022 Submitted (2022)
Keywords: disentanglement
dynamical systems
prediction
generative models
robustness
out-of-distribution
distribution shift
causal inference

Abstract: Deep neural networks have become increasingly of interest in dynamical system prediction, but out-of-distribution generalization and long-term stability still remains challenging. In this work, we treat the domain parameters of dynamical systems as factors of variation of the data generating process. By leveraging ideas from supervised disentanglement and causal factorization, we aim to separate the domain parameters from the dynamics in the latent space of generative models. In our experiments we model dynamics both in phase space and in video sequences and conduct rigorous OOD evaluations. Results indicate that disentangled models adapt better to domain parameters spaces that were not present in the training data while, at the same time, provide better long-term predictions in video sequences.

Under review as a conference paper at ICLR 2022
DISENTANGLED GENERATIVE MODELS FOR ROBUST DYNAMICAL SYSTEM PREDICTION
Anonymous authors Paper under double-blind review
ABSTRACT
Deep neural networks have become increasingly of interest in dynamical system prediction, but out-of-distribution generalization and long-term stability still re- mains challenging. In this work, we treat the domain parameters of dynamical systems factors of variation of the data generating process. By leveraging ideas from supervised disentanglement and causal factorization, we aim to separate the domain parameters from the dynamics in the latent space of generative mod- els. In our experiments we model dynamics both in phase space and in video sequences and conduct rigorous OOD evaluations Results indicate that disentan- gled VAEs adapt better to domain parameters spaces that were not present in the training data. At the same time, disentanglement can improve the long-tew and out-of-distribution predictions of faae--of-the-art models in video sequences 2
INTRODUCTION
The robust prediction of dynamical systems behaviour remains an open question in machine learning. and engineering in general. The ability to make robust predictions is important not only for forecasting systems of interest like weather (Garg et al. 2021 Ravuri et af.) 202T) but even more so because it enables innovations in fields like system control, autonomous planning (Hafner et al., (018)) and computer aided engincering (Brunton et al. 2020). In this context, the use deep generative models has recently gained significant traction for sequence modelling (Girin et al. 2020).
Robustness of machine learning models can be considered along two axes: long-term prediction and out-of-disttibuttin (OOD) performance. Accurate long-term prediction can be notoriously difficult in many dynamical systems, because error accumulation can diverge in finite time (Zhou et al. 2020: Raissi al. 2019), problem that even traditional solvers can suffer from. More importantly, machine learning techniques are known to suffer from poor OOD performance (Goyal & Bengio, 2020 i.e. when they employed in setting they had not encountered during the training phase.
Before addressing the OOD problem, we must first define what constitutes as OOD dynamical systems. We start by the observation that even simple dynamical systems, i.e the swinging pendulum or the n-body system, can have multiple continuous parameters that affect their evolution. These parameters can be manifested as differential equation coefficients, boundary or initial conditions etc. Our starting point is to consider distinct ranges of those parameters as separate domains. Under this view, becomes apparent why OOD prediction of dynamical systems can be hard: capturing the whole range of those parameters in single training set is unrealistic (Fotiadis et al. (2020) and further inductive biases are required (Miladinović et al. 2019: Bird & Williams/ 2019) Barber et al.,
2021). From dynamical system point of view, different parameters can produce widely different trajectories in phase space. motivating example can be bifurcations which occur when small change in the parameters of system causes sudden qualitative change in its behaviour.
We focus on the inductive bias disentangled representations for which the dynamics are separated from the domain parameters. Many approaches based on the use of neural networks try to jointly learn the dynamics and the physical parameters, which results in convoluted representations and usually leads to overfitting (Bengio et al. 2012). System identification can be used extract parameters, but
Coode reproducing our experiments at:
https://anonymoun hwwp://wwwuss.u......o... wwwp://wee....o.ee.sceee/eeee//e
dis-dyn -o-dyy-systeens
*Animated phase-space and video predictions are available https://bbtt.t.o -ysttmm

---

Under review as a conference paper at ICLR 2022
requires knowledge of the underlying system to be computationally effective (Ayyad et al. 2020). We, instead, leverage advances in Variational Autoencoders (VAEs) (Kingma & Welling 2014) that enable learning disentangled representations. Disentanglement enables different latent variables to focus on different factors of variation of the data distribution, and has been applied in the context of image generation (Higgins al., 2017 Kim & Mnih, 2018). This can be extended to modelling dynamical systems by looking at disentanglement from a causal perspective: from all the generative models which can have the same marginal distribution, identify the one with the true causal factors. To map this idea to sequence modelling we treat the domain parameters of a dynamical system as factors of variation. Recent findings (Locatello et al.] 2018) emphasize the vital role of inductive biases from models or data for useful disentanglement. Unsupervised disentanglement, based on the assumption of domain stationarity, a promising direction (Miladinović et al. 2019 Li & Mandt 2018). Nevertheless, this leaves a wealth of ground truth domain parameters, which can be cheaply collected in simulated datasets. This type of privileged information originating from simulations has been shown to be effective for domain adaptation in computer vision tasks (Sarafianos al. [017]] Lee et al,, 2018). We thus use supervised disentanglement (Locatello et al.) (019) by leveraging the ground truth domain parameters. To the best of our knowledge, using domain parameters information this way. has not been previously explored in the dynamical system prediction setting.
Contributions While others have treated domain parameters as factors of variation in the data distri- bution, our work is the first, to the best of our knowledge, that explicitly uses privileged information from simulated data to disentangle those domain parameters from dynamics in a supervised way. We furthermore conduct experiments both in the low-dimensional phase space of 3 dynamical systems and the nigh-dimensional video rendering of swinging pendulum. Disentanglement has, in the past, been mostly applied to VAEs because they are easily amenable to it. We additionally apply disentanglement on a more powerful, hybrid, model with both stochastic and deterministic parts (Hafner et al.) 2018). In doing so, we not only assess disentanglement on a generative model outside boundaries of VAËS but furthermore we do it on a model which is considered state-of-the-an in long-term video prediction (aaxeaa et al. 2021). In all cases, the prediction performance is assessed both in- in-distribution and also in OOD settings of increasing degrees of distribution shift. To our understanding, this is the first time such a rigorous OOD test performed. Our results in phase-space demonstrate that disentangled models can better capture the variability of dynamical systems com- pared to baseline models both in-distribution and OOD. In modelling dynamics in video sequences, results indicate that disentanglement is beneficial both for long-term prediction and OOD prediction.
Limitations This work focuses on dynamical system prediction. While the results can potentially open up many applications in general time-series modelling, this is out of the scope of this work. We prioritize to empirically study OOD downstream task performance and the inspection of the disentangled representations with appropriate metrics is left out of scope in this work.
2 RELATED WORK
VAEs and disentanglement Disentanglement aims to produce representations where separate factors of variation in the data are encoded into independent latent components. This be seen as finding the true causal model of the data. While supervised disentanglement in generative models is a long-standing idea (Mathieu al. (2016), ierffommationtthorreeee
properties can be leveraged allow unsupervised disentanglement in VAEs (Higgins et al. 2011 Kim & Mnih 2018)) The impossibility result from (Locatello al., 2018) suggested that disentangled learning is only possible by inductive biases coming either from the model or the data. Hence, the focus shifted back to semi- or weakly- supervised disentanglement approaches (Locatello et al., 2019 2020). While most of these methods focus on disentanglement metrics, we opt to directly assess using a downstream prediction task.
Disentanglement in sequence modelling While disentanglement techniques are mainly tested in a static setting, there is a growing interest in applying it to sequence dynamics. Using a bottleneck based on physical knowledge, Iten et al. (2018) learn an interpretable representation that requires conditioning the decoder time, but it can return physically inconsistent predictions in OOD data (Barber et al. (021). Deep state-space models (SSMs) have also employed techniques for disentangling content from dynamics (Fraccaro et al.) 2017 Li & Mandt 2018), ,, but, focus mostly on modelling variations in the content, failing to take dynamics into account. In hierarchical approaches
(Karl et al. 2017), different layers of latent variables correspond to different timescales: for example,
2

---

Under review as a conference paper at ICLR 2022
Supervised disentanglemmnt
Supervised disentanglement
Nost
Piurn Decodar
Encoder
Decoder
Mata
KL-DIvergance
loss
KL-Divergence
Prediction
loss
The VAE-SD model (Left), From an n-dimensional phase space input, a o-dimensional prediction of future time-steps is derived. The loss function has three parts: the reconstruction loss replaced by a prediction loss, the Kl.-divergence enforces the prior on to the latent space, and the extra loss term enforces the supervised disentanglement of domain parameters in latent space. The CNN-VAE-SD (iggh,) The input frames are first encoded to low-dimensional space, analogous phase space. Then, before the VAE-SD prediction scheme is applied recursively. The low-dimensional predictions are then decoded back to pixel space.
Figure 1:
in speech analysis for separating voice characteristics and phoneme-level attributes (Hsu et al. 2017). In approach similar work, Miladinović al. (2019) separate the dynamics from sequence- wide properties in dynamical systems like Lotka- Volterra, but do so in an unsupervised way which dismisses wealth of cheap information and only assesses OOD generalization in a limited way.
Feed-forward models for sequence modelling Deep SSM models are difficult to train as they require non-trivial inference schemes and careful design of the dynamic model (Krishnan al.. 2013] Karl
Karl et al. 2017). Feed-forward models, with necessary inductive biases, have been used for sequence modelling dynamical systems (Greydanus al. 2019 Fotiadis et al. 20200). In works like Hamiltonian Neural Networks Greydanus et al. (2019) the domain is fixed; together with Barber
al. (02T)) our work is an attempt tackling domain variability.
Privileged information for domain adaptation. Using privileged information during training has been shown to help with domain adaptation in computer vision tasks. Using segmentation masks of simulated urban scenery can improve semantic segmentation on the target domain (Lee et al. 2018). while clip art data can help with domain transfer in an action recognition task (Sarafianos et al. 2017).
3 METHODS
3.1 VARIATIONAL AUTOENCODERS
Variational autoencoders (VAEs) Kingma & Welling 2014) offer a principled approach to latent variable modeling by combining variational inference model q (z|æ) with a generative model (=) As in other approximate inference methods, the goal is to maximize the evidence lower bound (ELBO) over the data:
(z/x) Lo.0(x) (x/x) [log z)] DKu(qo(z (x)((((()))((())) x))lp(z))
(1)
The first part of the ELBO is the reconstruction loss (in our case the prediction loss) and the second part the Kullback-Leiller divergence that quantifies how close the posterior is to the prior.
Design choices for the model We use isotropic unit Gaussian prior 2) = N( 0, J) which
helps to disentangle the learned representation (Higgins al. 2017). The approximate posterior (encoder) distribution is Gaussian with diagonal covariance ( x) N (= M Ez), al- lowing a closed form KL-divergence, while the decoder has a Laplace distribution Po 2) Laplace (æ: p.. II) with constant diagonal covariance y 0, which is tuned empirically. This leads to an L1 loss that provides improved results in some problems (Mathieu et al. (018)) and empirically works better in our case. The parameters μ4 = p=( Ø). Σ. = diag lo_ (x; Φ) and μa = (2; computed via beed-forwand neural networks.
3

---

Under review as a conference paper at ICLR 2022
1.50
OOD Easy
Ensy hard
DOD Hard
1.45
Pondulum Length
Instial
1.20
300 135 170
Initial Angle
Gravity
Figure 2: Parameter distribution for the video pendulum test-sets. The initial angle and angular velocity are drawn from the same uniform distribution for all test-sets. For the in-distribution test-set we draw the pendulum length and gravity from the same distribution as during training. The OOD test-sets represent distribution shifts of increasing magnitude, where parameters are drawn from totally different space which has zero overlap with the training and in-distribution test-set.
3.2 DISENTANGLEMENT OF DOMAIN PARAMETERS IN LATENT SPACE
Apart from the disentanglement that stems from the choice of prior p(z), we explicitly disentangle part of latent space that it corresponds the domain parameters of each input sequence. We achieve this by using a regression loss term Le (Z <) between the ground truth factors of the domain parameters & E R and the output of the corresponding latents, Zi:k- We, empirically, opted for an L1 loss, corresponding methods have reported that binary cross-entropy works better than L2 (Locatello et al. (2019) but this does not fit well setting like ours. We hypothesize that BCE works better because of the implicit scaling. To address this, we propose applying a function G(Hz) which linearly scales the Pz between the min and max values of the corresponding factor variation:
Laplacian prior with mean & and unitary covariance. Previous
(Mz) Ma, (max(ε)) min(E)) min(E:)
(2)
where are the domain parameters and their corresponding minimum and maximum values of domain parameters from the training set. In all cases, the regression term is weighted by a parameter 8 which empirically tuned. Plugging these choices results in the following loss function:
Lo.0 =B, g0 (=)m))) |I (x(2; logy
(Prediction loss)
+ || (mimi 1) log diag (a= (mi
o)) (m)m ll (mim) 10 (KLe-Divegg)nc)
(3)
+6 +ōllex (Mzra (#imn |
(Sup. disentangl. loss)
Using the reparameterization trick (Kingma & Welling 2014), the loss is amenable to optimzation by stochastic gradient descent, with batch size n. The model architecture can be seen in Figure lefft)
3.3 DISENTANGLEMEN
VIDEO DYNAMICS
We further investigate the effect of disentanglement in video sequence dynamics. To this end, two generative models are used. The first is derived from the VAE formulation of the previous section and is called CNN-VAE and is similar to the VAE with the addition of a convolutional encoder and a decoder. The encoder projects the input frames down to low-dimensional space which can be thought equivalent the phase space of the system. A VAE is applied in this projection predict in the future coordinates the "phase space". The decoder then maps the predictions of the VAE back pixel space. The schematic of the model can be seen in Figure rright).
The second model we use the Recurrent State Space Model (RSSM) which has been successfully used for planning (Hafner al.) 2018). Since RSSM is hybrid model combining deterministic and variational components, it allows us to assess disentanglement outside the limited scope of VAEs. Furthermore, using fate-oofthhe-ae model in long-term video prediction (Sxxena et al. 202T). allows
4

---

Under review as a conference paper at ICLR 2022
Pendiulum
Lotka-Volterr
3-body system
8.5
1.5
Taabsot
003 Thak-st
CCD Tat-set
Tebbect
0O0 Thabact
DOD Tua-sed
Tasbast
000 Tet-wet
000 Thakast
VAE
VAE-SD
VAE-SSD
Figure 3: Disentanglement scaling methods. MAE for the first 200 time-steps. Boxplots show the top best performing sets of hyperparameters of each architecture. Both VAE-SD & VAE-SSD outpeform the VAE in all 3 systems. VAE-SSD better captures the parameter space of the original test-set but in most cases VAE-SI generalizes better OOD.
1e.1
1e.2
3-body system
2.5
3.2
7.0
15
6.5
6.0
Tet-eat
003 Tewkset Exsy
0O0
000 Tep-on
Taehent Hord
Hard
Exsy
Hard
Easy
MLP
MLP-SD
VAE
VAE-SD
Figure 4: Disentanglement in VAE vs MLP. MAE for the first 200 time-steps. Boxplots show the top 5 best performing sets of hyperparameters of each architecture. While disentanglement in VAE-SD consistently improves results, disentanglement in MLP-SD does not always generalize well OOD, producing unstable predictions for the OOD Lotka- Volterra datasets.
identify the limits of applying disentanglement in competitive models. The loss function shares the same formulation as in the original work of Hafner et al. (2018) with the addition of the supervised disentanglement loss. Since in the RSSM formulation there latent variables for each time-step, we apply a disentanglement loss on all of them, which empirically is set to be L2:
T
LRSSM-SD Σ
Balerlo lo [In p(o, 8x) Bg(er-loga-) [KL (p))) (--))]]
reconstruction
prediction
|| 8, [1
BBq(allos
(4)
supervised disentanglement
Where is the observations, the stochastic latent variables at time f, & are the k dimensional domain parameters and 6 tunes the supervised disentanglement strength.
EXPERIMENT - ODE PHASE SPACE DYNAMICS
DATASETS
In the phase-sa phase-space space experiments compare the models three well studied dynamical systems, the swinging pendulum, the Lotka- Volterra equations used to model prey- prey-predator populations, and the planar 3-body system:
5

---

Under review as a conference paper at ICLR 2022
Pendulum
Lotka Volteral
3-body system
3.0
Predator
0.0
2.0
-0.5
-1
-1.0
-2
-1.5
-1.0
-0.5
0.0
0.5
1.0
3
5
Prey
Ground truth
MLP
VAE
VAE-SD
x
Noisy input
Figure 5: Model predictions in phase space. Trajectories are taken from the OOD Test-Set Hard of each system. The model imput noisy. The circle and bold "x' markers denote the start and end of the ground truth trajectories respectively.
Simple pendulum: Ø+ sin 0 = x= Bxy
m; d%: di = K1 Σj
mp
rij
3-body system: de dt
F11
Lotka-Volterras ý = 6xy yy
K2V,, 1,2,3
The systems where chosen for varied complexity in terms of degrees freedom, number ODE equations and factors of variation. For the pendulum consider one factor of variation, its length I; Lotka-Volterre has 4 factors of variation a. B. y, 6 and the 3-body system has also 4 factors of variation K1, m,, m2, m3. Factors same between the training, validation and test sets. To further assess the OOD prediction accuracy, we create two additional test sets with factor values outside of the original range. We denote these datasets as OOD Test-set Easy and Hard, representing a smaller and bigger shift from the original range. As visual example, the distribution of the factors of variation for the Lotka- Voltera system is illustrated in Figure/oo the Appendix. The data were additionally corrupted with Gaussian noise. Dataset details can be found on Table of the Appendix.
drawn uniformly from a predetermined range which is the
4.2 MODELS AND TRAINING
The main goal of these experiments is to assess whether OOD prediction can be improved by disentangling dynamical system parameters in latent space of VAEs. We opt to use simple models to allow more experiments and comparisons. Our main baseline is the VAE upon which we propose two enhancements that leverage supervised disentanglement. The first VAE model, termed VAE-SD termed VAE-SSD uses an additional linear scaling function G ((Hz)) for the latent variable mean vector Mag as described in Section 3.2 Another baseline multilayer perceptron (MLP) autoencoder which allows comparison with a deterministic counterpart of the VAE. We also use supervised disentanglement on the latent neurons of the MLP, a model we refer MLP-SD. This enables us to assess if the privileged information can improve deterministic models. Lastly, we include an LSTM model, popular choice for low dimensional sequence modelling YYu et al. 2019). representative recurrent method.
supervised disentanglement without scaling function while the second model
Early experiments revealed significant variance on the performance of the models, depending on hyperparameters. Under these conditions, we took various steps to make model comparisons as fair as possible. Firstly, all models have similar capacity in terms of neuron count. Secondly, we tune various hyperparameter dimensions, some of which are shared, while others are model-specific. Third, we conduct a thorough grid search on the hyperparameters to avoid undermining model (details can be found in Tables 4l and 5 of the Appendix). Lastly, train the same number of experiments for all models which amounts to 1440 trained models in total, as summarized in Table 2 of the Appendix.

---

Under review as a conference paper at ICLR 2022
SSIM
Test-oat
1.00
OCD Tastaat Easy
OOD Test-sst Hard
155N
0.95
CNN-992.0
3.90
0.55
0.85
0.65
2.00
200
4.00
100
200
4.0
100
100
4.00 Timestop
100
Timestop
Timestop
PSNR
OOD Test Easy
Hard
30
K00M
3x
30
6AE-4D
2x
20
200
600
200
Timestop
Timestep
Timestop
Figure 7: Prediction quality on the video pendulum. SSIM(top) and PSNR(bottomg as a function of the distance predicted into the future (x axis)
4.3 RESULTS
For each dynamical system we focus on the performance on the three test-sets, the in-distribution test set, which shares the same parameter distribution with the training set, and the OOD test-sets (Easy and Hard), which represent an increasing parameter shift from the training data. Models are compared on the cumulative Mean Absolute Error (MAE) between prediction and ground truth for the first 200 time-steps. We consider this to be sufficiently long-term as it is at least 20 times longer than the prediction horizon used during training. Long predictions are obtained by re-feeding the model predictions back as input. This approach has been shown to work well in systems where the dynamics are locally deterministic (Fotiadis et al., 2020). A summary of the quantitative results can be found in Figures B & 4 and Table To account for the variability in the results, we present summary of the best runs of each model, selected by validation MAE. We generally observe that model performance is correlated to the distribution shift of test-sets, and this is consistent for all systems and models. The MAE is increasing as we move from the in-distribution test-set to the OOD Easy and Hard test-sets. Nevertheless, not all models suffer equally from the OOD performance drop.
Comparing the VAES (Figure 3). we see that disentangled VAE models offer a substantial and consistent improvement over the VAE across all dynamical systems. The improvement more pronounced for the OOD test-sets where the distribution shift is greater, a strong indication that disentanglement of domain parameters is an inductive bias that can lead better generalization. We also observe that VAE-SSD models the in-disribbution data better that VAE-SD. This seems to come slight overfitting cost, because the VAE-SD provides better OOD extrapolation in most cases. This could be explained because the scaling function is dependent factors of the training set. The extra information allows the model to better capture the training data but sacrifices some generalization capacity.
min and max values of the
On the other hand, disentanglement results for the MLP are mixed. While in-distribution MLP-SD offers better results than the plain MLP, the OOD test-sets, MLP-SD only performs favourably in the pendulum data. Furthermore in Lotka- Volterra, MLP-SD models are very unstable, and this is a drawback that affects some VAE-SD model too (see Table of the Appendix). Probabilistic models seem better suited to capture the variation in the data. The contrast between VAE-SD and MLP-SD illustrates that making use of privileged information and latent space disentanglement are not trivial
7

---

Under review as a conference paper at ICLR 2022
CNN-VAE
CNN-VAE-SO
RSSM
RSSM-SD
t-10
20
40
50
100
150
200
300
400
500
600
700
800
Figure 8: Absolute difference between ground truth and predictions on the test-set the pendulum data set.
and more work is needed to help us understand what works in practice and why. Lastly, the LSTM (Figure IT& & Tf the Appendix) is only comparable in the pendulum dataset and only for small OOD shifts. Qualitative predictions can be found in Figure
5 EXPERIMENT VIDEO SEQUENCE DYNAMICS
In the first experiment we assessed supervised disentanglement for phase space prediction, where the states the input trajectories are fully observable and only the domain parameters are unknown. This experiment extends the idea of supervised disentanglement to pixel-space input and output, where the physical states have to inferred by the model.
5.1 DATASETS
The dynamical system we use in this experiment the swinging pendulum, common benchmark for modelling dynamics in video sequences (Brunton et al. 2020: Barber et al., 202I). We consider 4 factors of variation, the length I, gravity and initial angle and angular velocity w. Factors are drawn uniformly from predetermined range. As before, we create test-set and two additional OOD test-sets (Easy and Hard). The OOD sets have length and gravity values outside of the original range, while the initial conditions 0, factors variation for the test-sets is illustrated in Figure 2 The trajectories are first computed in phase space using a numerical simulator and then rendered video frames of 64 x 64 pixels. More details about the dataset can be found in Section A.2 of the Appendix.
drawn from the same distribution. The distribution of the
5.2 MODELS AND TRAINING
In this experiment we use two different models CNN- and RSSM. ENN-VAE described in Section 3.3 and architectural details can be found in Section B.2.1 During training the CNN-VAE the inner VAE recursively used to predict, the number recursions being a hyperparameter (Table of the Appendix). We found that this type of training leads to more stable long term predictions. In total, 48 CNN- VNN-VAE models are trained half of which are with supervised disentanglement (CNN-VA) VAE-SD). The RSSM model is generative model including both stochastic and deterministic component. We only use supervised disentanglement on the stochastic part, and term that model RSSM-SD, Disentanglement applied all four factors of variation of the domain, despite only length and gravity varying between datasets. Detailed architectural and training details can be found in Section B.2.2 the Appendix.
5.3 RESULTS
Figure shows the quality of predictions video pendulum two metries: structural similarity (SSIM) and peak signal-to-noise ratio (PSNR) as a function of predicted time distance. We select the models which have the best cumulative metrics over the first 800 timesteps
validation set.
For the CNN- VAE, effects of disentanglement are more subtle. We observe that, in-distribution, the disentangled CNN- VAE-SD has very similar quality when compared the CNN-VAE. For the OOD
8

---

Under review as a conference paper at ICLR 2022
datasets, though, disentanglement offers improved long-term predictions. The improvement is more noticeable on the OOD test-sets, indicating that disentanglement can help with OOD robustness. For RSSM, we first note that both models perform significantly better than the CNN-VAE models, which is expected since they are considered competitive in long-term video prediction. Disentanglement in RSSM seems to produce trade-off. The plain RSSM model better in short-term prediction but its performance deteriorates with time, reaching VAE-CNN levels in all metrics. On the other hand, the RSSM-SD model provides the best long-term scores all metrics and all datasets. Qualitative results in Figure 5 show that almost all models produce accurate short time predictions (approximately up to 200 time-steps). This further stresses the importance of disentanglement for long-term performance. In terms of OOD robustness, disentanglement also appears to be helping. While the RSSM-SD model lacks in short-term prediction quality on the in-distribution test-set, this performance gap closes as the OOD test-sets get harder. More specifically, on the in-distribution test-set the RSSM-SD overtakes RSSM in SSIM after around 400 frames, while in the OOD Easy and Hard test sets, this happens around 350 and 250 time-steps respectively. This narrowing gap indicates robustness improves with increasing distribution shifts. The above findings are corroborated by LPIPS (Zhang et al.) 2018) comparisons (Figure [3 and Table of0 of the Appendix). Furthermore, the qualitative results show that all models accurately capture the appearance of the pendulum even long-term. Where they differ is on how well they capture the dynamics of the pendulum movement. This could offer an explanation why disentangling the domain from the dynamics is important, and why in practice offers better long- teggterm and out- -diistribttttnnn performance.
Overall, experiments suggest that supervised disentanglement can be used to model dynamical systems in video sequences, resulting improved long-term and OOD performance.
6 CONCLUSIONS
Using supervised disentanglement of domain parameters in generative models is promising avenue for improving robustness. Our experiments show that it can improve both OOD generalization and long-term prediction of dynamical systems. This was demonstrated in phase-space with VABs and also in video sequence modelling with state-of-the-art RSSMs.
By treating the domain parameters as factors of variation of the data and applying supervised disentanglement, several inductive biases are potentially enforced. First, the model in addition to prediction also performs "soft" system identification which acts as a regularizer. Second, it creates an implicit hierarchy such that some latent variables correspond to sequence-wide domain parameters and the rest capture instant dynamics. We speculate that this could additionally make the latent space more interpretable. Third, if the model can correctly extract the parameters this mean that the prediction based on both of them which is closer to how numerical integrators work, where the domain is known. All of these could lead the model to learn the correct causal structure of the data. Nevertheless, using privileged information for OOD robustness is not always straightforward and requires further exploration. This is evident by the results of the MLP autoencoders which do not yield as consistent improvements. criticism of our method could be that cheap privileged information is not always available and/or depends on using simulated data. Firstly, training on simulations is an increasingly appealing option because it is a cheap way to generate data to begin with. This is, also, clearly demonstrated by the many advancements (Peng et al. 2017) that try to bring models trained in simulated data to the real world. So there seems to no reason not to use the privileged information that comes with simulated data. Under that light supervised disentanglement can provide pathway for real world applications where robustness in dynamical system prediction is critical. Applying the method to other datasets where there are more complex dynamic can increase its relevance. Sequence-wiee wide parameters could also be exploited through self-supervision.
techniques like sim2real
REPRODUCIBILITY STATEMENT
We provide all the necessary code to reproduce experiments at the anonymous repo https 7/anonymous. 4open. science // httdddssdddsosoo
ems (will be de-anonymized after the review process). The repo contains code for generating all the datasets from scratch and also code for training all the models presented in this work. The README also contains instructions on how to train the models. The hyperparameters have used are clearly and thoroughly presented in the

---

Under review as a conference paper at ICLR 2022
Appendix. These steps should significantly help others reproduce experiments. For any further clarifications, you are encouraged to contact the corresponding author(s).
REFERENCES
Abdulla Ayyad, Mohamad Chehadeh, Mobammad I. Awad, and Yahya Zweiri. Real-time system identification using deep learning for linear processes with application to unmanned acrial vehicles. IEEE Access, 8: 1:122939- 122553, 2020. doi: 10. 00.1099AACCESS222000000000 2020.300627...
Gregory Barber, Mulugeta A. Haile, and Tzikang Chen. Joint Parameter Discovery and Generative Modeling of Dynamic Systems. 2021. htt.hhttpp//rxxi..orrrg/
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation Learning: A Review and New
Perspectives. IEEE Transactions on Pattern Analysis and Machine Intelligence, 35(8):1798-188) 2012. URL http:
org/ abs wetp/120..5500 5538
Alex Bird and Christopher K.I. Williams. Customizing sequence generation with multitask dynamical systems. arXiv, (i), 2019. ISSN 23318422.
Steven L. Brunton, Bernd R. Noack, and Petros Koumoutsakos. Machine Learning for Fluid Mechanics. Anmual Review of Fluid Mechanics, 52(1):477-508, 2020. ISSN 0066-4189, doi: 10.1 -044:aaaurr---uu--011199-9---000
Stathi Fotiadis, Eduardo Pignatelli, Mario Lino Valencia, Chris D Cantwell, Amos Storkey, and Anil A Bharath. Comparing recurrent and convolutional neural networks for predicting wave propagation. In ICLR 2020 Workshop on Deep Differential Equations, 2 2020. URL http:
wwwpsiivrrr/ab//// 2002. 08981
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, 201 (eeteecee 1):602-36111 2017. ISSN 10495258.
Sagar Garg, Stephan Rasp, and Nils Thuerey. Weatherbench probability: Medium-range weather forecasts with probabilistic machine learning methods. EGU General Assembly Conference Abstracts, pp. EGU21-11448, 2021.
Laurent Girin, Simon Leglaive, Xiaoyu Bie, Julien Diard, Thomas Hueber, and Xavier Alameda- Pineda. Dynamical Variational Autoencoders: A Comprehensive Review. 8 2020. URL http:
//aaxiiv.org wwt.:/orggabbaaaaeooooo /ts.///00 12595
Anirudh Goyal and Yoshua Bengio. Inductive Biases for Deep Learning of Higher-Level Cognition. 11 2020. URL http:://aaiiiii
abs http:/aab////111111 15091
Sam Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian Neural Networks. pp. 1-15, 2019.
URL htttthtt///xiiiii org. abs aww.aaass..... 1906. 01563
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning Latent Dynamics for Planning from Pixels. 36th International Conference Machine Learning, ICML 2019, 2019-June:4528 45577, 11 2018. URL http://arxiv.org
abs 1811.0455
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick. Mohamed Shakir, and Alexander Lerchner. beta-VAE: LEARNING BASIC VISUAL CONCEPTS WITH A CONSTRAINED VARIATIONAL FRAMEWORK. 44(6))807-831, 2017. ISSN 1078- 0874. doi: 10. 10.1177/107808840822280
Wei-Ning Hsu, Yu Zhang. and James Glass. Unsupervised Learning of Disentangled and Interpretable Representations from Sequential Data. Advances in Neural Information Processing Systems, 2017- 3icemberr1879-18855
2017. URL
077022
Raban Iten, Tony Metger, Henrik Wilming, Lidia del Rio, and Renato Renner. Discovering physical concepts with neural networks. 2018. http:h:tp///r///.//////////
10

---

Under review as a conference paper at ICLR 2022
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick Van Der Smagt. Deep variational Bayes filters: Unsupervised learning of state space models from raw data. Sth International Conference on Learning Representations, ICLR 2017 Conference Track Proceedings, (ii):1133, 2017.
Hyunjik Kim and Andriy Mnih. Disentangling by Factorising. 35th International Conference on Machine Learning. ICML 2018, 6:4153 4171, 2018. http:::paaaaww../////aaaa.
05983
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In 2nd International Conference on Learning Representations, ICLR 2014 Conference Track Proceedings. International Conference on Learning Representations, ICLR, 12 2014. URL httpp::/aariiv: wwtp:iiiiorg/aa/.. abs//
1312. 6114v10
Rahul G. Krishnan, Uri Shalit, and David Sontag. Deep Kalman Filters. 11 2015. URL http:
//axxiv.org/abs//5555 05121
Kuan-Hui Lee, Germán Ros, Jie Li, and Adrien Gaidon. SPIGAN, adversaiiall leaming from simulation. CoRR, abs/1 8 810.03756, 2018.
03756
Yingzhen Li and Stephan Mandt. Disentangled Sequential Autoencoder. 35th International Confer- ence on Machine Learning, ICML 2018, 33:8992-9001, 3 2018. URL htt.httpp//aaxxi..orttt
abs/ /003/100.... 02991
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf. and Olivier Bachem. Challenging Common Assumptions in the Unsupervised Learning of Dis- entangled Representations. 36th International Conference on Machine Learning. ICML 2019, 12359
2019-Jume:7247-728.. 11 2018. htt.t::p//aaxx...rggggpas/p
Francesco Locatello, Michael Tschannen, Stefan Bauer, Gunnar Rätsch, Bernhard Schölkopf. and Olivier Bachem. Disentangling Factors of Variation Using Few Labels. 5 2019. URL http:
//arxiv. http:/v...rr/bb//999000 htt.abbb/19900000 01258
Francesco Locatello, Ben Poole, Gunnar Rätsch, Berhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-SSuperviseed Disentangleemnn Without Compromises. arXivv 2 2020. URL http://arxiv.org/
02886
Emile Mathieu, Tom Rainforth, N. Siddharth, and Yee whye Teh. Disentangling Disentaggeeentnnnnn in 812. 811./3333
Variational Autoencoders. 12 2018.
Michael Mathieu, Junbo Zhao, Pablo Sprechmann, Aditya Ramesh, and Yann LeCun. Disentangling factors of variation in deep representations using adversarial training. 11 2016. URL http:
7/a www.ooogg/aaaaaaaoooooooo 611 03383
ore Miladinović, Muhammad Waleed Gondal, Bernhard Schölkopf, Joachim M. Buhmann, and Stefan Bauer. Disentangled State Space Representations. 6 2019. httptttt/////xi...tt
1906. 03255
Xue Bin Peng, Marcin Andrychowicz, Wojciech Zaremba, and Picter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. CoRR, abs/1710.00637, 2017. URL http: w/arxii.orr////o.o..a...... http://bs/1110..5.55
Maziar Raissi, P. Perdikaris, and G. E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differen- tial equations. Journal of Computational Physics, 378:686-700 707,, 2 2019, ISSN 10902716. doi: (0.1016//.jcp.22 2018..00.1......
Suman Ravuri, Karel Lenc, Matthew Willson, Dmitry Kangin, Remi Lam, Piotr Mirowski, Megan Fitzsimons, Maria Athanassiadou, Sheleem Kashem, Sam Madge, Rachel Prudden, Amol Mand- hane, Aidan Clark, Andrew Brock, Karen Simonyan, Raia Hadsell, Niall Robinson, Ellen Clancy, Alberto Arribas, and Shakir Mohamed. Skilful precipitation nowcasting using deep genera- tive models of radar. Nature, 597(7878):672-77,, 677, 2021. ISSN 1476-4687 doi: 10.1038/ ht..httsss: /tt.tttsst///oi.ioooooo..../ org 1038/ 586 021 03854-7

---

Under review as a conference paper at ICLR 2022
Nikolaos Sarafianos, Michalis Vrigkas, and Ioannis A Kakadiaris. Adaptive svm+: Learning with privileged information for domain adaptation. In Proceedings of the IEEE International Conference Computer Vision Workshops, pp. 2637-2644, 2017.
Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork Variational Autoencoders. 2021. URL wtg/abs////ooooo..oa..........rr abs/ 22022200222 09532
http://arxixs org.orggaa.
Yong Yu, Xiaosheng Si, Changhua Hu, and Jianxun Zhang. A Review of Recurrent Neural Networks: LSTM Cells and Network Architectures. Neural Computation, 3111 07 2019. ISSN
0899-7667.
doi: 10.1162/neco.a aat.cccooa_011111 199. URL
01199
Richard Yi Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, and Oliver Wang. The Unreasonable Effectiveness of Deep Features Perceptual Metric. Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recogmition, pp. 586-595, 12 2018, doi: 10.1 CVPR.2018.00065 URL /ttp/////////ppprraaaaaaaaa////r
Ehe hatp://cccccceeeetttteeeeeereerr
Haoyi Zhou, Shanghang Zhang. Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting. arXiv, 2020.
ISSN 23318422.
APPENDIX
A DATASETS
A.1 PHASE SPACE
For simulations, we use an adaptive Runge-Kute Kutta integrator with timestep of 0.01 seconds. Each simulated sequence has different combination of factors of variation. Simulation of the pendulum uses an initial angle 0 which is randomly between 10" 170" while the angular velocity w is O. For the other two systems the initial conditions are always the same avoid pathological configurations.
12

---

Under review as a conference paper at ICLR 2022
Table 1: Datasets. In L-V and 3-body OOD test sets, least domain parameter is outside of the parameter range used for training.
Pendulum
Lotka- Volterra
3-Body
mi di K1
ODEs
i $ sin 0 =
ax - -xy ÿ = 6xy - yy
d dr
K2
Number of ODEs Independent Variables
1 0,
2 x(prey), y(predator)
6
x (1.00-)),,
0 [10" - 170"]
Initial values
x x= 5
x = (1,-),, (0.5
x3 (0.1), (-0.5.0.0)
Timestep
0.01
0.01
0.01
Sequence length
2000
1000
1000
Noise 2
0.05
0.05
Factors of variation
(length)
a.B.y.
K2, m1 m2. My
(K2 2.055
A (ar c 2.055)
M = {ml 2.095 2.05]
B= [B e [0.95, 1.05]]
M2 {m2 [1.95, 2.05]]
Train/Val/Tess
lE [1.0 1.5]
C= 43.55
M3 {M3 [1.95,
(6 E 21.95
GOOD Hard
Qrain x (xxx)
(K M1 X M2 x M3)
A {a 2.094
K (K₂ [1. [1.94, 2.06]
B [B [0.94,
M {m] [1,94,.006
(y с [3.94,4.06])
{m2 [1.94, 2.06
OOD Test Set Easy
Ie [1.5 1.6]
D = [6 E 21.9.
M3 {m3 (1.94,
GOOD Easy = (A x B C
Qrain
ROOD Hard
(K X M X M2 M3) Quain
(a с 2.09],
K {K2 [1.93,2.07]]
B = [B E [0.93, 1.07])
M {m] 2
C (y é [3.93.4.0 4.07]]
{m² 1199, 2.07
OOD Test Set Hard
E [0.9 1.0]
D (6 [1.93, 2.07]}
M3 {mL 2.07]
GOOD Hard
ROOD Hard
(AxBxCxD) (Q ROOD Easy)
(K x M1 X M2 M3) (Quaain &2ODD Easy
Number of sequences Train/Val/Tes:
1000/1000/1000
OOD Test Set Easy
1000
OOD Test Set Hard
1000
13

---

Under review as a conference paper at ICLR 2022
1.07
Tesh-set
107
Testsset
000 mest Easy
1.6
COD rest set Easy
COD Hard
954
2.59
2.84
2.35
Figure 9: Example illustration of the parameter distribution for the LV sets. The regions do overlap, colors represent regions not boundaries. The OOD Easy test (orange) set does not include any of the parameter configurations of the training and original test set (green). Respectively, the OOD) Hard dataset (blue) does not include none of the OOD Easy the original test set configurations. The parameter space of the blue region is almost half as big at the green area (again without any overlap), signifying significant OOD shift).
14

---

Under review as a conference paper at ICLR 2022
A.2 VIDEO PENDULUM
This data set contains image sequences of a moving pendulum under different conditions. The positions of the pendulum first computed by numerical simulator and then rendered in pixel space as frames of dimension 64 x 64. An example image sequence shown in Figure [O] For the simulations, we use an adaptive Runge-Kute Kutte integrator with a timestep of 0.05 seconds. The length of the pendulum, the strength of gravity and the initial conditions (position, momentum) are set to different values so that each trajectory slightly differs from the others. The initial angle and initial velocity are drawn from the same uniform distribution for all data sets. The initial angle ranges from 30" to 170" and the initial velocity ranges from -2rad/s to 2rad/s s. For training, validation and in-distribution testing set, the gravity ranges from 8. 8.0m² /s 12.0m² /s, and the pendulum length ranges from 1.20m to 1.40m. the easy out-of-distribution testing set, the gravity is between 12.0 12.5m² and the pendulum length is between 1.40 1. 1.45m, while in the hard -iistrrbbitii testing set, the gravity is 12.5 13.0m") and the pendulum length is 1.45 1. 1.50m. The distributions of these parameters shown in Figure
I-10
t-11
t-12
1-15
1-14
t-15
t-16
t-17
t-10
I=19
Figure 10: Example image sequence from the video pendulum data set.
B TRAINING AND HYPERPARAMETERS
B. PHASE SPACE
During training the pack-propagation is used after single forward pass. The input and output the models are smaller than the sequence size, so cover the whole sequence we use random starting points per batch, both during training and testing. Both the VAE and MIP AE have an encoder with two hidden layers size 400,200 and a reverse decoder. The LSTM model has two stack LSTM cells with hidden size of 100, which results on an equivalent number of neurons. We used the Adam optimizer with b 0.9 and b2 0.999, A scheduler for the learning rate was applied whose patience and scaling factor are hyperparameters. Maximum number of epochs was set to 2000 but we employed also early stopping using validation which led to significantly less epochs.
Table 2: Number experiments with phase space data. Each experiment corresponds distinet configuration of hyperparameters.
MLP
MLP-SD
VAE
VAE-SD
VAE-SSD
LSTM
Total
Pendulum
72
72
72
72
72
72
432
L-V
72
72
72
72
432
3-body
96
96
96
96
96
96
576
Total experiments
1440
15

---

Under review as a conference paper at ICLR 2022
Table Pendulum hyperparameters
MLP
MLP-SD
VAE
VAE-SD
LSTM
Input Size
10, 50
Output Size
1, 10
Hidden Layers
[400, 200]
50,100,200
Latent Size
4, 8,
Nonlinearity
Leaky ReLU
Sigmoid 1,2,3
Num. Layers
Learning rate
10-3
Batch size
16, 32
16
16, 32
16
16, 64
Sched. patience
20, 30, 40
20,30
20
20
30
Sched. factor
0.3
0.3
0.3
0.3
0.3
Gradient clipping
1.0
1.0
Layer norm (latent)
No
No
Yes
Yes
No Partial
Teacher Forcing
Decoder y
10-3 10 10-5
10-3, 10-4
Sup. scaling
Linear
Linear
Supervision
0.1. 0.2. 0.3
0.01, 0.1. 0.2
# of experiments
72
72
72
72
72
Table 4: Lotka-Volterry hyperparameters
MLP
MLP-SD
VAE
VAE-SD
LSTM
Input Size
50
Output Size
10
Hidden Layers
[400, 200]
50,100
Latent Size
16,
Nonlinearity
Leaky ReLU
Sigmoid
Num. Layers
1,2,3
Learning rate
10-3 10
10-3, 10
10-3, 10-4
10-3
10
Batch size
16, 32, 64
16, 32
16, 32
16
10, 64, 128
Sched, patience
20, 30
20, 30
20
20
20, 30
Sched. factor
0.4 0.4
0.3
0.3
0.3
0.3
Gradient clipping
No
No
0.1,
0.1, 1.0 No
No
Layer norm (latent)
No
No
No
No
Teacher Forcing
Partial, No
Decoder
10-4 10-5 10-6
10-4 10-5, 10-6
Sup. scaling
Linear
Linear
Supervision 5
0.2. 0.3
0.01, 0.1. 0.2. 0.3
# experiments
72
72
72
72
72
16

---

Under review as a conference paper at ICLR 2022
Table 5: 3-body system hyperparameters
MLP
MLP-SD
VAE
VAE-SD
LSTM
Input Size
50
Output Size
10
Hidden Layers
[400, 200]
50,100
Size
8, 16, 32
Nonlinearity
Leaky ReLU
Sigmoid
Learning rate
10-3, 10
10-3, 10
10 10
10-3
Batch size
16, 32
16
16, 64, 128 20, 30
Sched. patience
30,40, 50, 60
30, 40, 50, 60
30, 40, 50, 60
30, 40, 50, 60
Sched. factor
0.3, 0.4
0.3
0.3, 0.4
0.3, 0.4
0.3
Gradient clipping
No
No
No
No
No
Layer norm (latent)
No
Decoder y Sup. scaling
10 1005,10 10
10-5, 10
Linear
Linear
Supervision
0.05, 0.1, 0.2,
0.1, 0.2
of experiments
96
96
96
96
96
B.2
VIDEO PENDULUM
B.2.1 CNN-VAE MODEL
Encoder has layers convolutional layers with 32, 32, 64 64 maps respectively. The filter size 3. padding is and stride is 2. The last convolutional layer is flattened 256-dimensional vector to become the inner VAE input. The decoder convolutional layers (64,64,32,32) with bi-linear upsampling. Model input and out is 20 frames. For the models without supervised disentanglement, grid search is performed upon the / value, the size of the latent space, and the roll-out length during training. For the models with supervised disentanglement, grid search is performed upon the value, the size the latent space, the time step of the data set, the roll-out length during training and the supervision multiplier. The detailed search grid is summarised in Table Learning rate was 10-3 and an Adam optimizer (b, = 0.9 and b2 0..99.99 0.999) was used. We also used early stopping upon the cumulative reconstruction loss for the first 200 steps on a validation set with the max number of epochs being 1000.
Table 6: Video pendulum hyperparameters for CNN-VAE models
CovnVAE
CovnVAE-SD
Latent Size
4, 8, 16 0.01, 0.1,1
8, 16
Decoder y VAE recursions
0.01, 0.
1,2, 4,
4, 8
Supervision
0.1,0.111
of experiments
36
36
B.2.2 RSSM MODELS
For the RSSM model we follow the architecture parameters described in Hafner et al. (2018 &
Saxena et al. (202T). For training we use sequences of 100 frames and batch size 100. All models were trained for 300 epochs with learning rate of 10-3 and Adam optimizer (b, = 0.9 and b2 0.999). During testing the model uses 50 frames as context (input). The parameters we tune appear in Table 7
17

---

Under review as a conference paper at ICLR 2022
Table 7: video pendulum hyperparameters for RSSM models
RSSM RSSM-SD
Batch Size Decoder std. Train Input Length Supervision 5 Seeds
50, 100
1.0, 2.0
50, 100
0.01, 0.1,
of experiments
24
C PHASE SPACE RESULTS
Pendiluum
Lotka-Nolterra
1e-2
3-body system
1.0
30
0.9
0.7
1.0
Easy
Hard
Easy
Hard
Basy
Hard
MLP-SD
VAE-SD
VAE-SSD
LSTM
Figure 11: Mean Absolute Error for first 200 time-steps. The bars represent the 5 experiments of each model with lower MAE. all three systems disentangled VAEs provide an advantage over other baseline. The disentanglement in MLP does not increase performance consistently. The scaling in VAE-SSD allows it to better capture the parameter space of the original test-set but in most cases VAE-SD generalizes better OOD, Similarly, LSTM does not generalize well OOD.
Table 8: MAE (x10²) of phase space experiments at 200 time-steps. Average and standard
deviation of best models.
Pendulum OOD Easy
Lotka- Volter
3-Body system
Test-set
OOD Hard
Test-set
OOD Easy
OOD Hard
Test-set
OOD Easy
OOD Hard
LSTM
11.16+0.70
19.49=1.06
30.9820.80
6.41=0.15
8.18+0.29
9.52=0.26
1.50 0.04
2.10 0.10
3.27=0.39
MLP
10.96=0.80
18.45=1.24
24.98=1.25
6.13=0.03
7.68=0.13
8.53=0.18
1.41 = 0.09
1.88 = 0.14
2.79+0.33
MLP-SD
10.08
16.61
25.50=0.92
6.07=0.03
N/A
N/A
1.32 a: 0.05
1.88 + 0.08
2.95+0.19
VAE
13.71
21.20=1.25
27.8642.48
6.22=0.04
1.84+0.13
8.66#0.03
1.39 + 0.04
1.83 + 0.06
2.9140.25
VAE-SD
10.92+0.28
16.40=1.75
23.62:2.09
6.05=0.05
7.46:0.07
8.16:0.21
1.35 + 0.02
1.70 + 0.07
2.27=1011
VAE-SSD
10.47+0.60
15.77=0.42
16.1
6.02=0002
7.56+0.14
8.38=0.09
1.32 0.04
1.74 0.03
2.27+0.12
Table 9: Models that diverge in at least trajectory. Percentage of the top 5 models selected by validation accuracy.
Pendulum
Lotka- Voltera
3-Body system
In Dist.
OOD Easy OOD Hard
In Dist.
OOD Easy
OOD Hard In Dist.
OOD Fasy
OOD Hard
LSTM
MIP
20%
MLP-SD
40%
100%
100%
VAE
VAE-SD
VAE-SSD
60%
18

---

Under review as a conference paper at ICLR 2022
Pendulum
Lotka Volterral
3-body system
1.0
3.0
0.5
Predator
2.0
-2
1.5
-3
5
Prey
Ground truth
MLP
VAE
VAE-SD
x Noisy input
Figure 12:
Model predictions in phase space. Trajectories are taken from the OOD Test-Set Hard of each system. The model imput noisy. The circle and bold 'x' markers denote the start and end of the ground truth trajectories respectively.
D VIDEO PENDULUM RESULTS
LPIPS
Easy
RSSM
RIDN-23 23
0.25
3.25
9.20
CRT Vie
2.20
0.15
3.10
3.30
3.05
3.00
0.00
200
4.00
600
500
200
4:0
600
800
200
4.00
600
800
Timestop
Timestop
Timestop
Figure 13: Prediction quality on the video pendulum as a function of the distance predicted into future (x axis). For LPIPS lower is better.
Table 10: Model comparison in video pendulum. Metrics calculated between ground truth and prediction the models at exactly 800 timesteps the future.
SSTM OOD Easy
I.PISS OOD Easy
PSNR OOD Easy
Test-set
OOD Hard
Test-set
OOD Hard
Test-set
OOD Hard
RSSM
0.795
0.787
0.783
0.227
0.246
0.252
13.36
12.71
12.26
RSSM-SD
0.813
0.808
0.794
0.192
0.206
0,233
14.16
13.82
12.90
CNN-VAE
0.787
0.781
0.784
0.242
0.259
0.250
12.67
12.41
12.36
CNN- VAE-SD
0.789
0.783
0.788
0.241
0.254
0.242
12.76
12.46
12.53
19

---

Under review as a conference paper at ICLR 2022
CNN-VAE
CNN-VAE-SD
RSSM
RSSM-SD
t-10
40
50
100
150
200
300
400
500
600
700
800
Figure 14: Absolute difference between ground truth and predictions on the OOD Test-set Hard of the video pendulum data set.
20


Reviews:
[{"review_id": "rcHyU-2sp1_", "reviewer": "Reviewer_dcJ9", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- A relevant problem in dynamical VAEs that is sufficiently motivated in the paper.\n- Empirical experiments on three problems: LV, video pendulum and the three-body problem. Demonstrate long term trajectory prediction and OOD on an easy and a hard task.\n- Have done a hyperparameter search and presented some ablation studies.\n\nI find it is an interesting work. However, I strongly feel the authors have not adequately demonstrated the benefit of disentanglement and are missing comparison. My main concerns are below:\n\n-  Lack of evidence on whether supervised loss disentangles dynamic parameters from domain parameters. The authors mention evaluation of disentanglement is beyond the scope. I beg to differ for two reasons: \n     - the long term trajectory prediction doesn’t necessarily benefit from the disentanglement of latent factors. There are several methods that achieve good performance in long-term trajectory prediction without any explicit form of disentanglement, for example, Hamiltonian Neural Networks, Hamiltonian generative network (HGN), Symplectic RNN [4], Physics-as-Inverse-Graphics [7], Lagrangian Neural Networks [6], etc.  In the related work section, the authors refer to HNNs and say disentanglement is not successfully addressed in such models. It would help if the authors could elaborate the sentence here. The HNNs learn Hamiltonians in a data-driven way and can make long term predictions. So why do they need to address disentanglement in the first place? It is unclear what added gain comes from supervised disentanglement and how it is advantageous over other state-of-the-art methods of long term trajectory predictions or OOD generalisation.\n\n     - Can supervised loss ensure the domain components are fully disentangled from dynamics? I think it is critical to demonstrate whether domain variables are disentangled from the dynamics in any meaningful way? Could, for instance, fix the domain variables and draw samples by slowly changing the dynamic variable and vice-versa. Or better report disentanglement metrics.\n\n- A weak baseline. As referenced above, several works on extending VAEs to dynamical models have shown empirical and theoretical (symplectic structure) arguments for long-term trajectory prediction. It would be worth comparing those methods and demonstrating any benefit in using a supervised approach. In addition, models like SINDY [5] can discover dynamical parameters in an unsupervised way and have demonstrated benefits on long-term trajectory prediction. Without comparison, it is not evident what the immediate benefits of the supervised setup are? If it is OOD generalisation, authors should at least show this as a limitation in existing approaches.\n\n- It is not apparent what makes the interval of domain parameters an easy or hard problem. It would be beneficial to discuss from the dynamical system perspective.\n\n- The choice of loss in supervised disentanglement needs more explanation. In Section 3.2, it is L1 and in Section 3.3 is L2.\n \n- In Table 9, the results of VAE-SD and VAE-SSD are unstable in some cases. But this is not the case with LSTM or VAE.  The authors should provide some discussion here and a potential explanation of the effect.\n\nMinor comments:\n\n- In Table 1, the domain parameters of the train/val/test set are in the same range. It is likely for a model to perform well on val/test if it has seen sequences of the same parameters in training. Shouldn’t the two be selected differently?\n- Please number the equations.\n\nTechnical inconsistencies:\n\n-In Section 3.2, the loss function is inconsistent with Figure 1. According to Figure 1, the input to VAE is x_n, and the prediction is x_{n+1}. The loss is a typical VAE plus a supervised disentangled term. There are no dynamics there. If the reconstruction term is supposed to be a prediction of x_{n+1} please add appropriate suffixes on x or z in $\\mu_x(z;\\theta)$. If this is not the case, please provide details on how dynamics are taken into account.\n\n- Please use consistent scripts. In Section 3.2, the k components of latent variable z are written as under script z_{1:k} and in Section 3.3 for latent variables as superscripts s^{1:k}.\n\n- In loss formulation of Section 3.2, the domain parameters $\\xi^{i}$ are associated with sample $x^{i}$. As far as I understand, the time steps in a sequence share the domain parameters. It would be helpful to use a suitable script to express it consistently.\n\n- In the loss formulation of Section 3.3, the prediction model is in the state space. The domain parameters are shared over T; why use the prediction model on $s_t$ instead of $d-k$ components of $s_t$? \n\n# References:\n\n[1] Chang MB, Ullman T, Torralba A, Tenenbaum JB. A compositional object-based approach to learning physical dynamics.\n\n[2] Sanchez-Gonzalez A, Bapst V, Cranmer K, Battaglia P. Hamiltonian graph networks with ode integrators. \n\n[3] Toth P, Rezende DJ, Jaegle A, Racanière S, Botev A, Higgins I. Hamiltonian generative networks. \n\n[4] Chen Z, Zhang J, Arjovsky M, Bottou L. Symplectic recurrent neural networks.\n\n[5] Champion K, Lusch B, Kutz JN, Brunton SL. Data-driven discovery of coordinates and governing equations.\n\n[6] Cranmer M, Greydanus S, Hoyer S, Battaglia P, Spergel D, Ho S. Lagrangian neural networks. \n\n[7] Jaques M, Burke M, Hospedales T. Physics-as-inverse-graphics: Joint unsupervised learning of objects and physics from video. \n", "summary_of_the_paper": "In this paper, the authors propose a supervised approach to disentangle domain parameters from the dynamics in a deep latent variable model like VAE. Extending VAEs to dynamical systems is a relevant problem and has been a focus of interest in many recent works [1,2,3,4,5,6,7]. This paper identifies two issues for developing dynamical VAEs, \n\ni) out of distribution generalisation\n\nii) long term trajectory prediction\n\nThe main contribution is to address the aforementioned issues using a supervised loss defined between latent variables and domain parameters. The authors present empirical experiments to support the idea.", "summary_of_the_review": "In my view, this paper proposes a fair approach to a relevant problem. However, there are several concerns. \n- The benefit of disentanglement is not demonstrated. The long term generation is not sufficient to support the claim. If a fully unsupervised approach can work equally good what is the incentive of supervised loss? Therefore, I think it is critical to compare with some of the methods outlined above.\n- The contribution is marginal as it simply introduces a regularisation term and provides empirical results. Simplicity is generally good and not a downside. However, it should be supported by proper justification and if possible perhaps by a theoretical claim. The choice of L1 in 3.2 and L2 in 3.3 is not properly explained.\n- There are technical inconsistencies that leave room for ambiguities.\n\nI have come to the conclusion this paper has concerns that need addressing. I, therefore, give a score of 5.\n\n# Post Rebuttal\n\nI have changed my score from 5 to 6.\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ict4xu8zV97", "reviewer": "Reviewer_2NUz", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- The novelty of this work may be not enough for ICLR acceptance standards in that the authors applied existing VAEs with minor modifications to known problems.\n- I think the comparison with unsupervised disentanglement models is quite unfair because the proposed model is trained using strong supervision about data. Furthermore, the results showing that supervised disentanglement methods outperform unsupervised ones are trivial and not particularly impressive. Instead of the baselines designed by the authors, it would be better to add some comparisons with existing papers for dynamical system prediction (particularly in Figures 3 and 4).\n- I am not sure whether the performance differences in Figures 3 and 4 are statistically significant because the results of some models exhibit quite high variances and the number of examined models (i.e., 5) seems small. It would be better to conduct some statistical tests to show that the differences are meaningful.\n- The experiments were conducted only on simple simulated datasets. I think some experiments on real-world and/or more complex data are necessary to show the applicability of the proposed method.\n- I think a deeper analysis on disentangle representations is not out-of-scope and is necessary because the paper heavily relied on VAEs for disentanglement learning. (i) Are the latent dimensions obtained with the supervision (z_1:k) truly disentangled? (ii) What kind of information is encoded in the other features without the supervision (z_k+1:d)? It would be better to add quantitative results based on existing disentanglement metrics and/or visual results (latent traversal, embedding space visualization).\n- Regarding Figure 7, it would be better to add proper explanation about why RSSM is better for the initial timesteps than RSSM-SD.\n- It would be better to improve the presentation quality of Figure 5. It is difficult to identify the differences between the lines in the current version because they are largely overlapped. Simply changing linear axis scales into log scales may be helpful.\n- Is the reconstruction loss in page 4 replaced by the prediction loss as described in the caption of Figure 1? If so, please modify the reconstruction loss in page 4 to accurately show the prediction loss.\n", "summary_of_the_paper": "The paper introduces a VAE-based disentanglement model for dynamical system prediction, which was trained under the supervision using domain parameters. The authors conducted experiments on simulated datasets and showed good performance for OOD cases and long-term predictions.", "summary_of_the_review": "I think the underlying technical contributions are quite small, while the empirical results are not particularly impressive. I thus find it difficult to argue for acceptance of the work.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "fu-9gwumYlm", "reviewer": "Reviewer_z3gj", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Contribution\n\nThe tackled problem of out-of-distribution generalization for the forecasting of dynamical systems is relevant and valuable to the community, as motivated in the paper. Models that generalize well in this setting still need to be discovered. This is an important issue as a forecasting system overfitting the training distribution cannot possibly have retained the true dynamics, in this case the ODE / PDE, of the observed phenomenon. The proposed method is one of the first steps in this challenging direction. However, to the best of my understanding, I find the contribution of this paper to be overly limited for acceptance, with numerous limitations that are described below.\n\n### Supervision and Disentanglement\n\nOne of the main claims of the paper is that the proposed method makes models disentangle the system's parameters, but this claim is not sufficiently supported. It is clear that the method makes model learn the system's factors of variations by design, but no experiment indicates that it does separate them as well, which is necessary for disentanglement. In this regard, I disagree with the statement at the end of Section 1 removing from the scope of the paper the inspection of the learned representations, given the claims of the paper.\n\nExperiments investigating the latent space could include, for example, the manipulation of the latent variables associated to the systems' parameters to assess whether they are actually disentangled, similarly to e.g. [1, Figure 1]. The advantage of the considered framework of experiments in this setting is that all the data is simulated: one could numerically compare the resulting sequence from the manipulation of the latent state with the actual simulated sequence generated with the corresponding parameters, in line with [2, Sections 5.2 and 5.3] in another context of disentanglement.\n\nRegardless of the disentanglement property, an experiment evaluating the ability of the models learned with the proposed loss to retrieve the true parameters of the system from the observation of the sequence would greatly emphasize the utility of the method. Indeed, estimating the parameters of dynamical systems is an active line of research which this paper directly follows, given that the method is supervised on these parameters.\n\n### Novelty and Significance\n\nThe novelty of the proposed method is limited, given prior works on supervised disentanglement in other contexts (e.g. [3] cited in the paper). The discovery that the supervision of the system's parameters improves forecasting performance is interesting but expected; an opposite result may have been questionable. I stress that this is not a significant problem per se, given that evidencing this behavior in the context of dynamical systems is valuable to the community. However, this limited novelty is to be considered with the lack of significance in the presented empirical study.\n\nThe lack of significance first lies in the numerical results, which all show non-significant or marginal improvements for the proposed method against the baselines: VAE-SSD performs similarly to VAE-SD, and VAE-SD is only marginally better than VAE and is far from closing the gap between out-of-distribution and in-distribution performance. Furthermore, qualitative animated results that I checked in the Hard OOD setting provided as supplementary do not match the example of Figure 5 as VAE and VAE-SD seem very close compared to the difference between them and the ground truth, making the improvement look thin. To my understanding, these mixed results may be the consequence of choices in experimental design, as argued in the next point of this review.\n\nThe experiments also lack significance in their design. The considered out-of-distribution parameter ranges are restricted and close to the training parameters. Without further discussion in the paper to contextualize this choice, it would seem that out-of-distribution sequences might be close to in-distribution sequences, thus questioning the obtained results. I would suggest the authors to either extend the considered ranges of parameters or explain how the current parameters are relevant.\n\nA possible direction to improve this paper in this regard might be to consider a semi-supervised setting, like [3], instead of a fully supervised method like in the current version. Real-world simulations can be expensive to run and the available labeled data in this context may be limited, thus rather motivating a method working with sparsely labeled data.\n\n### Choice of Models\n\nThe choice of models to apply the proposed supervision may be questionable and explain the mixed results obtained in the experiments. To the best of my knowledge, forecasting models of the kind of VAE and CNN-VAE in the paper are not widely adopted in the community; I would be interested in references that the authors could provide to support this choice. Instead, state-of-the-art variational models often rely on sequential latent variable generative models, like [4, 5, 6, 7, 8] to name some works in the last five years. Moreover, the use of ODE-like recurrent schemes may be considered as well as they have been shown to be adapted to the prediction of dynamical systems (see e.g. [9, 10, 11]).\n\nRSSM is the only baseline in this paper following this line of works, and is also the only model presenting a substantial improvement with the proposed supervision. I believe that this is no coincidence, given that a questionable choice of model like MLP may induce suboptimal results with the introduced method. I encourage the authors to strengthen their empirical evaluation by considering more robust and standard models.\n\n\n## Other Remarks and Questions\n\n### Questionable Claims\n\nSeveral other claims is the paper may be questionable, as listed in the following.\n\n> System identification [...] requires knowledge of the underlying system to be computationally effective\". [Page 1]\n\nIt would seem that the proposed method does require knowledge of the underlying system as well, since it relies on supervising over the system's parameters.\n\n> [We treat] the ground truth domain parameters from simulations as privileged information which, to the best of our knowledge, has not been applied to dynamical system prediction previously. [Page 2]\n\nThis may be a wording issue but privileged information has already been leveraged for dynamical systems for the last few years, cf. for example [10, 11, 12], even though this privileged information is not necessarily the system's parameters. The authors might consider further discussing this point.\n\n> The problem is that [VAEs] usually lack in competitive performance.\n\nWithout references to support this claim, I would strongly disagree given the references mentioned above [4, 5, 6, 7, 8].\n\n### Number of Experiments\n\nFigures 3 and 4 are said to show the top 5 models of each architecture, but I could not understand the details of this selection. Does this correspond to the top 5 best performing sets of hyperparameters? Or is it the top 5 over a given number of experiments for the same set of hyperparameters?\n\n### LPIPS\n\nCould the authors justify the choice of LPIPS for the experiments in Section 5? LPIPS is a perceptual metric for realistic images, making it a priori less relevant for synthetic datasets like these pendulum sequences. The authors might rather highlight PSNR which is a standard metric for this type of datasets and is already used in the appendix.\n\n### Writing\n\nThe paper is mostly clear and easy to read, but I find the description of the models to be confusing regarding their nature and the considered architectures (for instance, the VAE is underspecified in the main text), which raises issues in the motivation of the modeling choices in the paper as mentioned above. Many figures are hard to read in greyscale; I recommend that the authors improve their readability to make them as accessible as possible.\n\nTypos:\n - the reference to Saxena et al. (2021) at the end of Section 1 should be between parentheses;\n - title of Section 3.2: \"disentanglment\" should be \"disentanglement\";\n - title of Section 3.3: there is an extra parenthesis at the end of the title;\n - Section 3.3: \"which can be though as\" should be \"which can be thought of as\";\n - there is an extra comma and no parentheses are needed in the last sentence of page 4;\n - Section 4.3: \"We also observe that VAE-SSD model the in-distribution data\" should be \"We also observe that VAE-SSD models the in-distribution data\".\n\n\n## References\n\n[1] I. Higgins et al. $\\beta$-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. ICLR 2017.\\\n[2] J. Donà et al. PDE-Driven Spatiotemporal Disentanglement. ICLR 2021.\\\n[3] F. Locatello et al. Disentangling Factors of Variations Using Few Labels. ICLR 2020.\\\n[4] R. G. Krishnan et al. Structured Inference Networks for Nonlinear State Space Models. AAAI 2017.\\\n[5] M. Fraccaro et al. A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning. NIPS 2017.\\\n[6] E. Denton et al. Stochastic Video Generation with a Learned Prior. ICML 2018.\\\n[7] J. Chung et al. A Recurrent Latent Variable Model for Sequential Data. NIPS 2015.\\\n[8] Y. Rubanova et al. Latent Ordinary Differential Equations for Irregularly-Sampled Time Series. NeurIPS 2019.\\\n[9] R. T. Q. Chen et al. Neural Ordinary Differential Equations. NeurIPS 2018.\\\n[10] S. Greydanus et al. Hamiltonian Neural Networks. NeurIPS 2019.\\\n[11] Y. Yin et al. Augmenting Physical Models with Deep Networks for Complex Dynamics Forecasting. ICLR 2021.\\\n[12] M. Raissi et al. Physics-Informed Neural Networks: A Deep Learning Framework For Solving Forward and Inverse Problems Involving Nonlinear Partial Differential Equations. Journal of Computational Physics. 2019.", "summary_of_the_paper": "This paper introduces a supervised loss to encourage dynamical systems predictors to retain systems' parameters (e.g. appearing in the underlying ODE) in their latent space. It presents multiple experiments to evaluate the advantages of this approach, including better long-term forecasting ability as well as improved prediction performance for out-of-distribution parameters, that had not been seen during training.", "summary_of_the_review": "From the limitations that are described in this review, I think that this paper needs very significant changes to be accepted, especially because of the questionable claims and insufficient experimental results. Nonetheless, I am looking forward to discussing my opinion with the authors and other reviewers. I believe that this paper follows an interesting line of research and that further work could make it ready for publication at a next conference.\n\n### Post-Rebuttal Update\n\nI acknowledge the authors' response and thank them for their extensive answer. As explained in my follow-up response, I find that the proposed improvements are marginal and insufficient to raise my score. Therefore, I maintain my strong recommendation to reject the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "N_5As25hiO7", "reviewer": "Reviewer_yFXa", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths:**\n* Clear statement of the underlying hypothesis being tested\n* Clear presentation of the results and supporting information\n* Extensive sweeps over hyperparameters\n\n**Weaknesses:**\n* Improvement of the models with disentanglement in the phase space setting appear marginal; Based on the provided visualizations it is not clear that there is a systematic way in which models with disentanglement perform better. A more expressive analysis of the errors might be helpful to assess this aspect (maybe distribution of errors across the dataset for several fixed samples?)\n* It’s hard to assess how much variance in the performance is present in the video prediction metric; This is general challenge with selecting best performing models, as they completely mask away the error bars; (Providing several model instances would help to evaluate the significance better)\n* While marginal improvements are presented in coarse performance metrics, an insight into the type/class of errors that are being reduced would be very interesting.\n* One potentially important hyperparameter (time step) was not varied, which often significantly affects the prediction accuracy.\n", "summary_of_the_paper": "The paper studies the performance of dynamical systems learned from data with a focus on out of distribution (OOD) evaluations. Authors consider the question whether disentangling dynamical system parameters in the latent space can improve the generalization of the models, which is perceived as privileged information available from the reference (ground truth) simulations. Authors carry out experiments on several dynamical systems: pendulum, Lotka-Volterra system and three-body problem. Additionally an experiment on video prediction of a singing pendulum is performed. Authors found that additional disentanglement can improve generalization performance of the models and in video prediction setting leads to better long-term predictions based on structural and perceptual image metrics.", "summary_of_the_review": "Authors present a clear investigation of how disentanglement of the domain factors may affect the performance of learned dynamical models. The suggested experimental evaluation is sound, but current results seem a little marginal. With additional results/modifications I believe this work could be useful to a wider audience, but my initial rating is marginally below the threshold. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "0cbxrB0TSM", "reviewer": "Reviewer_HYc5", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- My main issue is the limited novelty of the proposed method. This method is a straightforward extension of the unsupervised disentangled state-space model (Miladinovic et al.) to a supervised one, where the privileged information regarding domain parameters is explicitly fed to the model.\n- Certain claims (e.g. regarding disentanglement) are made without proper quantitative and/or qualitative investigation(s).\n- It is claimed that the proposed supervised disentanglement method improves performance over the unsupervised method. However, there are no comparisons with the closest unsupervised method  (e.g. Miladinovic et al.). Therefore it is hard to judge whether the proposed supervised method truly performs better or not.\n- The results on OOD generalization can not be considered OOD as the parameters' range used to create OOD datasets highly overlap with the ranges of the training datasets (Table 1, Appendix).\n\nI expand on these points in the following:\n\n- Regarding line 1 in the contributions section: The treatment of domain parameters as factors of variation is one of the main proposals of DSSM in Miladinovic et al. DSSM seeks to disentangle these true parameters (also referred to as domain-invariant state dynamics) from observations only. Therefore, I believe it's not the first work to consider this setting.\n- The main contribution of this work is the supervised disentanglement of sequential data. However, the authors do not investigate how good the disentanglement is. This leaves room for interpretation i.e. whether it is truly disentanglement that is helping the model in achieving good performance. I believe both quantitative and qualitative analysis of disentanglement would further strengthen the claims made in this paper.\n- The true parameters (factors of variation) are explicitly provided to the network for training. The method is not directly comparable to Locatello et al. 2019, as only a few labels were used in Locatello et al. 2019 which resulted in a semi-supervised disentanglement setting, in contrast to a fully supervised setting in this work.\n- The authors claim that the supervised disentanglement of the sequential model is better than the unsupervised disentanglement done in DSSM  (Miladinovic et al.). I would appreciate it if the authors could back it up with some empirical evidence. It is important to compare results with DSSM (even Kalman VAE) to see the true benefits of supervision.\n- The method is practically limited as the privileged true parameter information is not readily available in real-world systems. Thus, as acknowledged by the authors, this method can only work for the simulated systems where these variables are known beforehand.\n- The ranges of the parameters used to create the OOD dataset highly overlaps with the ranges used to create the training dataset. In my opinion, this is not OOD as it is very likely that the test sample comes from the range which is used for training. I suggest authors use the ranges which are completely outside the ranges of the training distribution i.e. extrapolation generalization regime (or even interpolation regime where the parameters are sampled from the subset of the training range but that subset range is not seen during training).\n- Have the authors tried Gaussian distribution (correspondingly L2 loss) for the decoder? I wonder how the results might differ from the Laplace distribution.\n- The prediction quality is reported by using perceptual metrics LPIPS and SSIM. These metrics compare the deep feature space and statistical properties of the images respectively. I think these metrics are not sufficient for evaluating the predictions of the dynamic. If it is possible then kindly report RMSE and/or NLL.\n- Fig1 caption: Do the input-output dimensions differ? I don't think the labeling in the figure is correct as there are some inconsistencies. For e.g. in a single time step: $x_1$ → $x_{n+1}$ and $x_n$ → $x_{n+o}$.\n\nTypos (minor):\n\npg2: \"This can be extend to\"\n\npg2:  \"high-dimemnsional video rendering\"\n\npg2:  The sentence \"we directly assess using the downstream prediction task.\" seems incomplete.\n\npg4: \"though as equivalent to the phase space of the system.\"\n\npg4: check sentence structure of \"being a state-of-the-art model in long-term video prediction,\"\n\npg5: \"on three well studies dynamical systems,\"\n\npg9: \"prediction is based both on them which\"", "summary_of_the_paper": "This paper introduces a supervised disentanglement method to learn dynamical systems. The method relies on the provision of privileged information (true parameters of a sequence) in order to disentangle them from observations. The method is evaluated on three toy datasets.", "summary_of_the_review": "I have some reservations regarding the novelty of this work. Moreover, certain empirical results do not back the claims made in this paper. Therefore, I am inclined towards rejection.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This manuscript tackles an interesting and significant line of research of long-term prediction and out-of-distribution generalization in time series models. I strongly believe this problem is an important one to solve. However, in its current form, its novelty is marginal, and the experiments fail to decisively show advantages. It also lacks of systematic improvements and error analysis. Further work could make it ready for publication at a next conference.
Decision: Reject