Identifying latent state transitions in non-linear dynamical systems

## Abstract

Abstract This work aims to improve generalization and interpretability of dynamical systems by recovering the underlying lower-dimensional latent states and their time evolutions. Previous work on disentangled representation learning within the realm of dynamical systems focused on the latent states, possibly with linear transition approximations. As such, they cannot identify nonlinear transition dynamics, and hence fail to reliably predict complex future behavior.
Inspired by the advances in nonlinear Ica , we propose a state-space modeling framework in which we can identify not just the latent states but also the unknown transition function that maps the past states to the present.
We introduce a practical algorithm based on variational auto-encoders and empirically demonstrate in realistic synthetic settings that we can (i) recover latent state dynamics with high accuracy, (ii) correspondingly achieve high future prediction accuracy, and (iii) adapt fast to new environments.

## 1 Introduction

We focus on the problem of understanding the underlying states of a target dynamical system from its low-level, high-dimensional sensory measurements. This task is prevalent across various fields, including reinforcement learning [Hafner et al., ] and robotics [Levine et al., ]. As a running example of such a system, we consider a drone controlled by an autonomous system. Here, the observational data would be a video stream (Fig. 1 (b)) instead of the full state of the system comprised of absolute position, velocity, and acceleration in $3D$ (Fig. 1 (a)). This system may have additional variables influencing the state evolution, e.g., the strength and direction of the wind at any time or drone motor torque set by the controller. The main objective of this work is to learn latent representations and state transition functions that would be useful for downstream tasks, e.g., computing the control signals for the optimal transport of the drone from point A to point B.

#### Identifiability.

Due to the partially observed nature of these problems, learning dynamics directly in the data space (e.g., pixel space) is not feasible, and previous works often focus on learning *latent dynamical systems* [Hafner et al., ]. However, such latent models commonly are not guaranteed to recover the true underlying states and transitions (*non- identifiability*), which results in entangled representations, lack of generalization across new domains, and poor interpretability [Schmidhuber, , Bengio et al., ].
Identifiable representation learning aims to address these challenges by learning the underlying factors of variation in the true generative model [Hyvärinen and Pajunen, ]. Most existing methods [Hyvarinen and Morioka, , , Hyvarinen et al., , Khemakhem et al., , Klindt et al., ] assume mutually independent components that do not affect each other. This is unrealistic for dynamical systems as the present state of the system depends on the past states, i.e., the transition function propagates the system state by (nonlinearly) mixing the past state components [Morioka et al., , Yao et al., , ].

#### Identifiability in dynamical systems.

Recently, Yao et al. [, ] showed that under certain assumptions, it is possible to *identify* or recover the true unobserved latent states in a dynamical system (up to component-wise transformations).
Klindt et al. [] show that temporally sparse transitions could recover the true latent variables up to trivial transformations. Morioka et al. [] introduce a framework to estimate the ‘innovations’ or process noise, which represents the stochastic impulses fed to a dynamical process.

Figure: Figure 1: Sketch of our method and main theoretical contribution. (a) We assume an underlying unobserved dynamical system, e.g., a cartpole, where the full state is composed of the cart position and velocity, and the angle and angular velocity of the pole: $[x,\dot{x},\theta,\dot{\theta}]$. (b) We partially observe the system as a sequence of video frames, which are used as input to our method. (c) We learn an inverse generative function that maps the raw observation signals to the latent state variables, as well as a transition function that maps the past latent states to the present latent state. Identifiability of the latent states is ensured by Theorem 1 [Yao et al., ]. In addition to this, our main contribution is the identifiability of the transition function ensured by Theorem 2.
Refer to caption: x1.png

Yao et al. [] recover latent states and identify their temporal relations from sequential data by utilizing non-stationarity noise. Later, Yao et al. [] extend their framework to instead exploit the autocorrelation between the past and the present latent states, while including factors modulating the dynamics and generative functions.
While these attempts lead to provably identifiable representations, they only propose non-parametric or linear approximations to the unobserved state transition function. While their nonparametric approximations cannot be unrolled over time, a linear model falls short of predicting future states of complicated systems (as we demonstrate empirically in our experiments).

#### Our contributions.

We present the first framework that allows for the identification of the unknown transition function alongside latent states and the generative function (see Fig. 1).
Following the previous works [Klindt et al., , Yao et al., , ], we first establish the identifiability of the latent states (Fig. 1: Theorem 1).
Different from these works, our framework allows estimating the process noise, representing the random impulses fed to a dynamical system. Inspired by Morioka et al. [], we show that the estimation of the correct transition function is ensured by restricting the process noise and the transition function (Fig. 1: Theorem 2). These restrictions form our main assumptions for identifiability: (i) the process noise shows non-stationarity or autocorrelation driven by an auxiliary variable and (ii) a trivially augmented transition function is bijective.

In our formulation, the process noise is encoded from observations similar to Franceschi et al. [].
While previous works perform backward prediction of noise from present and past states [Yao et al., , ], we show that encoding the process noise and predicting future latent states from the past latent states and process noises is the key to achieving an identifiable transition function.
We propose a straightforward evidence lower bound that allows us to recover true underlying factors in the limit of infinite data.
Our empirical findings show that our framework manages to predict the future states of an unknown system and is easier to adapt to novel environments.

## 2 An identifiable dynamical system framework

This section starts with the notation and our generative dynamical model that leads to identifiable variables and functions under certain assumptions. Next, we state the assumptions to achieve identifiability, followed by our main theoretical contribution.

### 2.1 Our generative model

We are interested in inferring latent dynamical systems from high-dimensional sensory observations $\mathbf{x}_{1:T}$, where $t$ is the time index and $\mathbf{x}_{t}\in\mathbb{R}^{D}$. As usual, we assume a sequence of latent states $\mathbf{z}_{1:T}$, with $\mathbf{z}_{t}\in\mathbb{R}^{K}$, are instantaneously mapped to observations via a generative function ${\mathbf{g}}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{D}$:

$$ $\mathbf{x}_{t}={\mathbf{g}}(\mathbf{z}_{t}).$ $$

Without loss of generality, the latent states $\mathbf{z}_{1:T}$ evolve according to Markovian dynamics:

$$ $\mathbf{z}_{t}=\mathbf{f}(\mathbf{z}_{t-1},\mathbf{s}_{t}),$ $$

where $\mathbf{f}:\mathbb{R}^{2K}\rightarrow\mathbb{R}^{K}$ is an auto-regressive transition function and $\mathbf{s}_{t}\in\mathbb{R}^{K}$ corresponds to additional variables influencing the dynamics, e.g., random forces acting on the system or control signals.

#### Augmented dynamics

For the ease of notation in the rest of the paper, we introduce the following generative process with augmented transition and generative functions [Morioka et al., ]:

$$ $\displaystyle\mathbf{z}_{0}$ $\displaystyle\sim p_{\mathbf{z}_{0}}(\mathbf{z}_{0}),$ # initial state (1) $\displaystyle\mathbf{s}_{t}$ $\displaystyle\sim p_{\mathbf{s}|\mathbf{u}}(\mathbf{s}_{t}|\mathbf{u})=\prod_{ k}p_{s_{k}|\mathbf{u}}(s_{kt}|\mathbf{u}),\quad\forall t\in 1,\ldots,T,$ # process noise (2) $\displaystyle\begin{bmatrix}\mathbf{z}_{t}\\ \mathbf{z}_{t-1}\end{bmatrix}$ $\displaystyle={\mathbf{f}_{\texttt{aug}}}\left(\begin{bmatrix}\mathbf{s}_{t}\\ \mathbf{z}_{t-1}\end{bmatrix}\right)=\begin{bmatrix}\mathbf{f}(\mathbf{z}_{t-1 },\mathbf{s}_{t})\\ \mathbf{z}_{t-1}\end{bmatrix},\quad\forall t\in 1,\ldots,T,$ # state transition (3) $\displaystyle\begin{bmatrix}\mathbf{x}_{t}\\ \mathbf{x}_{t-1}\end{bmatrix}$ $\displaystyle={{\mathbf{g}}_{\texttt{aug}}}\left(\begin{bmatrix}\mathbf{z}_{t} \\ \mathbf{z}_{t-1}\end{bmatrix}\right)=\begin{bmatrix}{\mathbf{g}}(\mathbf{z}_{t })\\ {\mathbf{g}}(\mathbf{z}_{t-1})\end{bmatrix},\qquad\forall t\in 2,\ldots,T.$ # observation mapping (4) $$

Above, $\mathbf{u}$ is an auxiliary variable modulating the noise distribution $p_{\mathbf{s}|\mathbf{u}}$. Recent nonlinear independent component analysis (Ica) works have achieved strong identifiability results by exploiting such auxiliary variables to obtain non-stationarities [Hyvarinen and Morioka, ] and autocorrelations [Hyvarinen and Morioka, ].
Inspired by these works, we consider the following cases:

- •
Setting $\mathbf{u}$ to an observed regime index leads to a nonstationary process noise. A real-world example would be a flying drone under different precipitation conditions, which can be observed up to a noise level.
- •
Setting $\mathbf{u}=\mathbf{s}_{t-1}$ implies an autocorrelated noise process. A real-world example would be a flying drone in a windy environment where the wind speed or direction changes continuously.

### 2.2 Identifiability theory

Let $\mathcal{M}=({\mathbf{f}_{\texttt{aug}}},{{\mathbf{g}}_{\texttt{aug}}},p_{
\mathbf{s}|\mathbf{u}})$ denote the ground-truth model. We learn a model $\hat{\mathcal{M}}=({\hat{\mathbf{f}}_{\texttt{aug}}},{\hat{{\mathbf{g}}}_{
\texttt{aug}}},\hat{p}_{\mathbf{s}|\mathbf{u}})$ by fitting the observed sequences. We make the following assumptions:

- (A0)
Distribution matching [Klindt et al., , Yao et al., , ] The learned and the ground-truth observation densities match everywhere:
$\displaystyle p_{{\mathbf{f}_{\texttt{aug}}},{{\mathbf{g}}_{\texttt{aug}}},p_{
\mathbf{s}|\mathbf{u}}}(\{\mathbf{x}_{t}\}_{t=1}^{T})=p_{{\hat{\mathbf{f}}_{
\texttt{aug}}},{\hat{{\mathbf{g}}}_{\texttt{aug}}},\hat{p}_{\mathbf{s}|\mathbf
{u}}}(\{\mathbf{x}_{t}\}_{t=1}^{T}),\forall\mathbf{x}_{t}\in\mathcal{X}.$
(5)
This is a standard assumption in the nonlinear Ica literature. It ensures that the model $\hat{\mathcal{M}}$ is sufficiently flexible, e.g., a neural network, that it learns the correct distribution in the limit of infinite data.
- (A1)
Injectivity and bijectivity [Morioka et al., ] The generator functions ${\mathbf{g}}$ and ${\hat{{\mathbf{g}}}}$ are injective, which implies that the augmented generative functions ${{\mathbf{g}}_{\texttt{aug}}},{\hat{{\mathbf{g}}}_{\texttt{aug}}}$ are injective. The augmented dynamics functions ${\mathbf{f}_{\texttt{aug}}},{\hat{\mathbf{f}}_{\texttt{aug}}}$ are bijective.
In contrast to Morioka et al. [], which use an augmented transition model on observations $(\mathbf{x}_{t-1},\mathbf{x}_{t})$, our formulation captures the functional dependence between a latent pair $(\mathbf{z}_{t-1},\mathbf{z}_{t})$ and the process noise $\mathbf{s}_{t}$.
- (A2)
Decomposed transitions [Klindt et al., , Yao et al., , , Song et al., ] Each dimension of the transition function $\{f_{k}\}_{k=1}^{K}$ is influenced by a single process noise variable $s_{kt}$. The output is a single latent variable $z_{kt}$:
$\displaystyle z_{kt}$
$\displaystyle=f_{k}(\mathbf{z}_{t-1},s_{kt}),\quad\text{for }k\in 1,\ldots,K
\text{ and }t\in 1,\ldots,T.$
(6)
- (A3)
Conditionally independent noise. Let $q_{k}(s_{kt},\mathbf{u})=\log p(s_{kt}|\mathbf{u})$ denote the conditional density of the noise variable $s_{kt}$. Let $\eta_{k}(z_{kt},\mathbf{u})=\log p(z_{kt}|\mathbf{z}_{t-1},\mathbf{u})$ denote the conditional density of the state variable $z_{kt}$. Conditioned on the auxiliary variable $\mathbf{u}$, we assume:
•
Each noise variable $\mathbf{s}_{t}\in\mathbb{R}^{K}$ is independent over its dimensions $s_{1t},\ldots,s_{Kt}$:
$\log p(\mathbf{s}_{t}|\mathbf{u})=\sum_{k=1}^{K}\underbrace{\log p(s_{kt}|
\mathbf{u})}_{q_{k}(s_{kt},\mathbf{u})}=\sum_{k=1}^{K}q_{k}(s_{kt},\mathbf{u})
,\quad\forall t\in 1,\ldots,T,$
(7)
•
the past latent state $\mathbf{z}_{t-1}$ and the present noise $\mathbf{s}_{t}$ are independent:
$\mathbf{s}_{t}\perp\!\!\!\perp\mathbf{z}_{t-1}|\mathbf{u}$.
Since $\mathbf{s}_{t}\perp\!\!\!\perp\mathbf{z}_{t-1}|\mathbf{u}$ and each dimension of the transition function $f_{k}(\mathbf{z}_{t-1},s_{kt})$ is a function of only a single (conditionally) independent noise variable $s_{kt}$, the conditional density of the latent pair $(\mathbf{z}_{t},\mathbf{z}_{t-1})$ also factorizes:
$\displaystyle\log p(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{u})=\sum_{k=1}^{K}
\underbrace{\log p(z_{kt}|\mathbf{z}_{t-1},\mathbf{u})}_{\eta_{k}(z_{kt},
\mathbf{u})}=\sum_{k=1}^{K}\eta_{k}(z_{kt},\mathbf{u}).$
(8)
The same is assumed for the learned conditional density: $\log p(\hat{\mathbf{z}}_{t}|\hat{\mathbf{z}}_{t-1},\mathbf{u})=\sum_{k=1}^{K}
\log p(\hat{z}_{kt}|\hat{\mathbf{z}}_{t-1},\mathbf{u}).$
This assumption generalizes the nonstationary noise assumption in Yao et al. [] which consider the more specific case of setting the auxiliary variable $\mathbf{u}$ to the regime index.
- (A4)
Sufficient variability of latent state $\mathbf{z}_{t}$. For any $\mathbf{z}_{t}$, there exist some $2K$ values of $\mathbf{u}$: $\mathbf{u}_{1},\ldots,\mathbf{u}_{2K}$, such that the $2K$ vectors
$\mathbf{v}(\mathbf{z}_{t},\mathbf{u}_{1}),\ldots,\mathbf{v}(\mathbf{z}_{t},
\mathbf{u}_{2K})$
(9)
are linearly independent for some index $l$ of the auxiliary variable $\mathbf{u}$, where
$\displaystyle\mathbf{v}(\mathbf{z}_{t},\mathbf{u})=\left(\frac{\partial^{2}
\eta_{1}(z_{1t},\mathbf{u})}{\partial z_{1t}\partial u_{l}},\cdots,\frac{
\partial^{2}\eta_{K}(z_{Kt},\mathbf{u})}{\partial z_{Kt}\partial u_{l}},\frac{
\partial^{3}\eta_{1}(z_{1t},\mathbf{u})}{\partial z^{2}_{1t}\partial u_{l}},
\cdots,\frac{\partial^{3}\eta_{K}(z_{Kt},\mathbf{u})}{\partial z^{2}_{Kt}
\partial u_{l}}\right)\in\mathbb{R}^{2K}.$
(10)
This assumption is a generalisation of the sufficient variability assumption in Yao et al. [] where the conditional density $\eta_{k}(z_{kt},\mathbf{u})$ conditions on a specific regime index rather than a general auxiliary variable.
It implies that the dependence of the latent states $\mathbf{z}_{t}$ on the auxiliary variable $\mathbf{u}$ is not too ’simple’, meaning that the corresponding conditional log densities $\eta_{k}(z_{kt},\mathbf{u})$ change with respect to $\mathbf{u}$ sufficiently differently for each value of $\mathbf{u}_{1},\ldots,\mathbf{u}_{2K}$, e.g., eliminating ’simple’ rotation-invariant Gaussian densities [Hyvarinen et al., , Hyvärinen et al., ].
- (A5)
Sufficient variability of process noise $\mathbf{s}_{t}$. For any $\mathbf{s}_{t}$, there exist some $2K$ values of $\mathbf{u}$: $\mathbf{u}_{1},\ldots,\mathbf{u}_{2K}$, such that the $2K$ vectors
$\mathbf{w}(\mathbf{s}_{t},\mathbf{u}_{1}),\ldots,\mathbf{w}(\mathbf{s}_{t},
\mathbf{u}_{2K})$
(11)
are linearly independent for some index $l$ of the auxiliary variable $\mathbf{u}$, where
$\displaystyle\mathbf{w}(\mathbf{s}_{t},\mathbf{u})=\left(\frac{\partial^{2}q_{
1}(s_{1t},\mathbf{u})}{\partial s_{1t}\partial u_{l}},\cdots,\frac{\partial^{2
}q_{K}(s_{Kt},\mathbf{u})}{\partial s_{Kt}\partial u_{l}},\frac{\partial^{3}q_
{1}(s_{1t},\mathbf{u})}{\partial s^{2}_{1t}\partial u_{l}},\cdots,\frac{
\partial^{3}q_{K}(s_{Kt},\mathbf{u})}{\partial s^{2}_{Kt}\partial u_{l}}\right
)\in\mathbb{R}^{2K}.$
(12)
Similar to (A3), this assumption ensures that the noise density contains enough structure in the form of autocorrelation or non-stationarity, where the process noise density changes with respect to $\mathbf{u}$ sufficiently differently for each value of $\mathbf{u}_{1},\ldots,\mathbf{u}_{2K}$.

#### Remark

Yao et al. [] used nonstationarity of the process noise for identifiability of the conditionally independent latent states (Theorem 1). If the variable $\mathbf{u}$ is an observed categorical variable (e.g., domain indicator), the assumptions (A4, A5) can be written in an alternative form without partial derivatives with respect to $u_{l}$ [Hyvarinen et al., ] (see Section A.5 for the alternative version).

### 2.3 Main theoretical contribution

In this section, we state our main theoretical contribution, that is, the identifiability result for the dynamical function $\mathbf{f}$ (Theorem 2). For completeness, we start with a theorem on the identifiability result for the conditionally independent latent states $\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{u}$ (Theorem 1), which Yao et al. [] established for the nonstationary noise case. The proofs are detailed in Sections A.3 and A.4.

#### Theorem 1 (Identifiability of latent states 𝐳 1 : T subscript 𝐳 : 1 𝑇 \mathbf{z}_{1:T} bold_z start_POSTSUBSCRIPT 1 : italic_T end_POSTSUBSCRIPT ) [Yao et al., 2021 ] :

Under assumptions (A0, A1, A2, A3, A4), latent states $\mathbf{z}_{t}={\mathbf{h}}(\hat{\mathbf{z}}_{t})$ are identifiable up to a function composition ${\mathbf{h}}=\pi_{z}\circ r_{z}$ of a permutation $\pi_{z}:[K]\rightarrow[K]$ and element-wise invertible transformation $r_{z}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{K}$. Or equivalently, the same follows for the generative function ${\mathbf{g}}\circ{\mathbf{h}}={\hat{{\mathbf{g}}}}$.

#### Theorem 2 (Identifiability of the dynamical function 𝐟 𝐟 \mathbf{f} bold_f ):

Under assumptions (A0, A1, A2, A3, A4, A5), the process noise $\mathbf{s}_{t}={\mathbf{k}}(\hat{\mathbf{s}}_{t})$ is identifiable up to a function composition ${\mathbf{k}}=\pi_{s}\circ r_{s}$ of a permutation $\pi_{s}:[K]\rightarrow[K]$ and an element-wise invertible transformation $r_{s}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{K}$. Equivalently, the dynamical function ${\mathbf{h}}^{-1}_{\texttt{aug}}\circ{\mathbf{f}_{\texttt{aug}}}\circ{{\mathbf
{k}}_{\texttt{aug}}}={\hat{\mathbf{f}}_{\texttt{aug}}}$ is identifiable up to a function composition ${{\mathbf{k}}_{\texttt{aug}}}=[{\mathbf{k}},{\mathbf{h}}]=\pi\circ r$ of a permutation $\pi=[\pi_{s},\pi_{z}]:[2K]\rightarrow[2K]$ and an element-wise invertible transformation $r=[r_{s},r_{z}]:\mathbb{R}^{2K}\rightarrow\mathbb{R}^{2K}$, where Theorem 1 already proves that ${\mathbf{h}}^{-1}_{\texttt{aug}}=[{\mathbf{h}}^{-1},{\mathbf{h}}^{-1}]$ is an invertible element-wise transformation.

## 3 Practical implementation using variational inference

We turn our theoretical framework into a practically usable implementation using variational inference, which approximates the true posterior over the noise and the states given the observations. For space considerations, we defer the details of our inference setup to Appendix B. Please see the below algorithm for implementation details and Figure 7 for architecture details.

Figure: Algorithm 1 Practical learning algorithm

#### Initial encoding.

We map each observation $\mathbf{x}_{t}$ to an initial embedding $\mathbf{r}_{t}$ via an Mlp or Cnn depending on the input modality. Using the first $T_{\text{ic}}=4$ initial embeddings $\mathbf{r}_{1:T_{\text{ic}}}$, an initial condition encoder (Mlp) outputs the parameters of the variational posterior $q(\mathbf{z}_{0}|\mathbf{x}_{1:T_{\text{ic}}})$ for the initial condition $\mathbf{z}_{0}$. Our ablations showed that the model is rather insensitive to $T_{\text{ic}}$.

#### Sequence prediction.

We start by sampling from the initial value distribution $\mathbf{z}_{0}\sim q(\mathbf{z}_{0}|\mathbf{x}_{1:T_{\text{ic}}})$.
Next, a forward sequential layer (Rnn + Mlp) takes the initial embeddings $\mathbf{r}_{1:t}$ up to time $t$ and the (sampled) previous latent state $\mathbf{z}_{t-1}$ as input $[\mathbf{r}_{1:t},\mathbf{z}_{t-1}]$, and outputs the parameters of the variational posterior $q(\mathbf{s}_{t}|\mathbf{x}_{1:t},\mathbf{z}_{t-1})$ for the noise variables $\mathbf{s}_{1:T}$. For example, for the first noise variable $\mathbf{s}_{1}$, the variational posterior is of the form $q(\mathbf{s}_{1}|\mathbf{x}_{1},\mathbf{z}_{0})$.
Subsequently, given a sample from the noise variable $\mathbf{s}_{1}\sim q(\mathbf{s}_{1}|\mathbf{z}_{0},\mathbf{x}_{1})$ and the sampled initial state $\mathbf{z}_{0}$, we predict the next state $\mathbf{z}_{1}\equiv\mathbf{f}(\mathbf{z}_{0},\mathbf{s}_{1})$, or more specifically: $z_{k1}\equiv f_{k}(\mathbf{z}_{0},s_{k1})$ for $k\in 1,\ldots,K$. We model each output $k$ of the transition function $f_{k}$ as a separate Mlp, to encourage the conditional independence of the latent states. We recursively compute the trajectory $\mathbf{z}_{2:T}$ of future latent states and noise posteriors $q(\mathbf{s}_{2:T})$.

#### Priors and ELBO computation.

We assume a standard Gaussian prior for the initial condition $p(\mathbf{z}_{0})=\mathcal{N}(0,I)$. The prior $p(\mathbf{s}_{t}|\mathbf{u})=\prod_{k}p(s_{kt}|\mathbf{u})$ over the noise variables are $1D$ trainable conditional flows. To allow for multi-modal prior distributions for $1D$ noise variables $s_{kt}$, we choose neural spline flows as the flow architecture [Durkan et al., , Stimper et al., ]. For the computation of the KL divergence, the terms containing a conditional flow do not have closed form solutions. We compute them by a Monte Carlo approximation using the sequence samples $\{\mathbf{z}_{0:T},\mathbf{s}_{1:T}\}$. Finally, the decoder $\mathbf{d}$ outputs the mean of our Gaussian observation model with fixed variance: $p(\mathbf{x}_{t}|\mathbf{z}_{t})=\mathcal{N}(\mathbf{x}_{t};\mathbf{d}(\mathbf
{z}_{t}),I)$.

## 4 Experiments

In this section, we perform experiments on two temporal datasets where the temporal dependencies are governed by an underlying latent dynamical system: (i) a synthetic dataset containing multivariate time-series data similar to [Hyvarinen and Morioka, , , Hyvarinen et al., , Yao et al., , ], and (ii) a more challenging reinforcement learning environment composed of high-dimensional video sequences of a physical cartpole system [Huang et al., , Yao et al., ]. Our goal is to answer the following research questions:

- (Q1)
Can our model identify the underlying latent dynamics, which contains three components: latent states $\mathbf{z}_{1:T}$, the transition function $\mathbf{f}$, the process noise $\mathbf{s}_{1:T}$?,
- (Q2)
What is the advantage of our model which jointly learns the latent variables and the transition function, compared to a plug-in method which first identifies the latent variables and simply fits a transition function on those?
- (Q3)
Ablation: How does each model component contribute to the disentanglement performance?
- (Q4)
By learning an identifiable transition function $\mathbf{f}$, can we predict the system state for longer horizons than the training sequence length?
- (Q5)
Is identifiable dynamics useful for sample efficient adaptation of dynamical systems?

#### Metrics.

We evaluate the performance on the latent dynamics identification using mean correlation coefficient (Mcc), the standard metric in the nonlinear Ica literature [Hyvarinen et al., , Yao et al., ].
Let us denote the latent states and the process noise for first $T_{\text{train}}$ steps by $\bar{\mathbf{z}}_{\text{train}}=\mathbf{z}_{0:T_{\text{train}}}$ and $\bar{\mathbf{s}}_{\text{train}}=\mathbf{s}_{1:T_{\text{train}}}$ respectively.
For (Q1,Q3), we measure the validation Mcc for the latent sequence $\bar{\mathbf{z}}_{\text{train}}$: $\textsc{Mcc}[\bar{\mathbf{z}}_{\text{train}}]$. For the synthetic data set, we additionally measure the validation Mcc for the noise sequence $\bar{\mathbf{s}}_{\text{train}}$: $\textsc{Mcc}[\bar{\mathbf{s}}_{\text{train}}]$.
For (Q2,Q4,Q5), we measure the future prediction performance.
For this, we take the subsequent $T_{\text{future}}$ steps coming after the first $T_{\text{train}}$ steps.
We denote the future observations and latent states for these frames as $\bar{\mathbf{x}}_{\text{future}}=\mathbf{x}_{T_{\text{train}}+1:}$ and $\bar{\mathbf{z}}_{\text{future}}=\mathbf{z}_{T_{\text{train}}+1:}$ respectively.
As metrics, we measure the mean squared error (Mse) on the future observations $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}}]$ and the Mcc on the future latent states $\textsc{Mcc}[\bar{\mathbf{z}}_{\text{future}}]$.
When $T_{\text{future}}$ takes different values, e.g., $T_{\text{future}}{}\in\{2,4,8\}$, we denote the future metrics by $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}[2]}]$, $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}[4]}]$ and $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}[8]}]$.

#### Baseline methods.

For (Q1), we compare with nonlinear Ica methods assuming: (i) independent latent processes $\mathbf{z}$ with no temporal or nonstationary structure: $\beta$-Vae [Higgins et al., ], (ii) mutually independent latent processes $z_{kt}=f_{k}(z_{k,t-1},s_{kt})$ with nonstationarity: Tcl [Hyvarinen and Morioka, ] and iVae [Khemakhem et al., ], (iii) mutually independent latent processes $z_{kt}=f_{k}(z_{k,t-1},s_{kt})$ with temporal structure: Pcl [Hyvarinen and Morioka, ] and SlowVae [Klindt et al., ], and (iv) temporally mixed latent processes $z_{kt}=\mathbf{f}(\mathbf{z}_{t-1},s_{kt})$ with nonstationarity: two versions of Leap [Yao et al., ] having linear and nonparametric transition functions Leap-Lin and Leap-Np. For (Q2), we construct a baseline based on the closest methods to ours: Leap-Np. We first identify the latent variables using Leap-Np, and then fit an Mlp on them as the transition function in an offline manner. For (Q4,Q5), we compare with a disentangled deep state-space model KalmanVae [Fraccaro et al., ], a nonlinear Ica method Leap-Lin with linear transitions that can predict the future states, and the constructed baseline Leap-Np-Mlp.

### 4.1 Synthetic experiment

**Table 1: Synthetic experiment results (mean std.dev. across 5 runs). For methods that cannot predict the future, we leave the $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}[\cdot]}]$ rows empty (N/A). Likewise, we compute $\textsc{Mcc}[\bar{\mathbf{s}}_{\text{train}}]$ only for Leap and our method as others do not maintain process noise variables.**
|  | Models |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Metrics | $\beta$-Vae | Pcl | Tcl | iVae | SlowVae | KalmanVae | Leap-Lin | Leap-Np | Ours |
| $\textsc{Mcc}[\bar{\mathbf{z}}_{\text{train}}]$ $\uparrow$ | $0.60$ | $0.57$ | $0.39$ | $0.58$ | $0.41$ | $0.64$ | $0.68$ | $0.89$ | $0.95$ |
| ($\pm 0.05$) | ($\pm 0.05$) | ($\pm 0.07$) | ($\pm 0.06$) | ($\pm 0.05$) | ($\pm 0.05$) | ($\pm 0.03$) | ($\pm 0.04$) | ($\pm 0.08$) |  |
| $\textsc{Mcc}[\bar{\mathbf{s}}_{\text{train}}]$ $\uparrow$ | N/A | N/A | N/A | N/A | N/A | N/A | $0.14$ | $0.26$ | $0.66$ |
| ($\pm 0.01$) | ($\pm 0.04$) | ($\pm 0.09$) |  |  |  |  |  |  |  |
| $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}[2]}]$ $\downarrow$ | N/A | N/A | N/A | N/A | N/A | $1.27$ | $0.22$ | N/A | $0.06$ |
| ($\pm 0.19$) | ($\pm 0.03$) | ($\pm 0.01$) |  |  |  |  |  |  |  |
| $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}[4]}]$ $\downarrow$ | N/A | N/A | N/A | N/A | N/A | $1.32$ | $0.18$ | N/A | $0.09$ |
| ($\pm 0.27$) | ($\pm 0.03$) | ($\pm 0.01$) |  |  |  |  |  |  |  |
| $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}[8]}]$ $\downarrow$ | N/A | N/A | N/A | N/A | N/A | $1.72$ | $0.59$ | N/A | $0.21$ |
| ($\pm 0.82$) | ($\pm 0.13$) | ($\pm 0.03$) |  |  |  |  |  |  |  |

Figure: Figure 2: Model predictions in the data space with the estimated uncertainties. Our model uses the first $T_{0}=2$ data points to encode the initial latent state $\mathbf{z}_{0}$ and $T_{\text{dyn}}=4$ data points to encode the noise sequence $\bar{\mathbf{s}}_{\text{train}}$. We unroll our model for $T_{0}+T_{\text{dyn}}+T_{\text{future}}=14$ steps ahead, where the future noise variables $\mathbf{s}_{7:14}$ follow the learned prior flow. We draw 32 trajectory samples $\mathbf{x}_{0:14}$ by sampling from the initial state and process noise. Above, black and blue dots show the training and test data points. The red curves are the mean trajectories and the red region corresponds to $\pm 2$ standard deviation computed empirically. We observe near-perfect predictions and low uncertainty for the input data (the first $T_{\text{train}}=T_{0}+T_{\text{dyn}}=6$ time points) while the uncertainty grows as we unroll over time. Further, the uncertainty grows even more when the model predictions are off. Therefore, almost all test points lie in the $\pm 2$ std region, reflecting the high calibration level our probabilistic model attains. Please see Figure 5 for all regimes.
Refer to caption: x2.png

As common in nonlinear Ica literature [Hyvarinen et al., , Morioka et al., , Yao et al., , ], we first set up a synthetic experiment to show that (Q1) our model recovers the latent dynamics, and hence (Q4) achieves a higher future prediction accuracy. On the same synthetic setup, we (Q3) analyze how each model component affects the performance in an ablation study, and show that (Q2) our method predicts the future behavior better than a plug-in method which trains a transition function on top of the identified latent states.

#### Dataset.

Similar to Yao et al. [, ], we set up a synthetic data experiment containing multivariate time-series. We set the dimension of $\mathbf{s}$, $\mathbf{z}$ and $\mathbf{x}$ to $K=8$. Same as Yao et al. [, ], we use a 2-linear layer random MLP as the generative function ${\mathbf{g}}$, 2-linear layer random MLP as the transition function $\mathbf{f}$ and we choose $\mathbf{z}_{1:T}$ to be a second-order Markov process, i.e., $\mathbf{z}_{t}=\mathbf{f}(\mathbf{z}_{t-2},\mathbf{z}_{t-1},\mathbf{s}_{t})$. The number of environments is $R=20$. For each environment, we generate 7500/750/750 sequences as train/validation/test data following our generative model. The distribution of the process noise $\mathbf{s}_{t}$ is conditioned on the environment index $\mathbf{u}$. Each sequence has length $T=T_{0}+T_{\text{dyn}}+T_{\text{future}}=2+4+8=14$. As we have a second-order Markov process, first $T_{0}=2$ states are spared as initial states. The next 4 observations $\mathbf{x}_{1:4}$ are used for training the dynamical model. The last 8 observations $\mathbf{x}_{5:12}$ are used for assessing the performance of future estimation. We choose the future prediction horizon $T_{\text{future}}=8$ as the double of the training sequence length. If the dynamics are truly identified, the model should predict future states well, even for a longer horizon.

#### Main results.

$\textsc{Mcc}[\bar{\mathbf{z}}_{\text{train}}]$, $\textsc{Mcc}[\bar{\mathbf{s}}_{\text{train}}]$ and $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}}]$ results are shown in Table 1. Our model recovers latent states $\mathbf{z}$ and process noise $\mathbf{s}$ better than the baselines, as demonstrated by a higher correlation with the true latent states and the true process noise (in terms of $\textsc{Mcc}[\bar{\mathbf{z}}_{\text{train}}]$ and $\textsc{Mcc}[\bar{\mathbf{s}}_{\text{train}}]$). This leads to a higher accuracy in future prediction, in terms of $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}}]$ with prediction horizons $\{2,4,8\}$. Notice that 8-step future prediction corresponds to the double the amount of dynamics steps the models see during the training. As the prediction horizon increases, we see that the difference in the future prediction performance between our method and the baselines also increases.

Figure: Figure 3: Mcc vs. Mse results for ablations for different time steps.
Refer to caption: x4.png

#### Ablation study.

To understand how each model component affects the identification and prediction performance (Q3), we set up an ablation study. We chose $\beta$-Vae [Higgins et al., ] as the base model. On top of it, we add three model components one by one: (i) $1$-Mlp-Dyn: a single Mlp modeling $K$ outputs of the transition function $\{f_{k}\}_{k=1}^{K}$, (ii) $K$-Mlp-Dyn: $K$ independent Mlps modeling $K$ outputs of the transition function $\{f_{k}\}_{k=1}^{K}$, and (iii) Ours: an additional conditional normalizing flow to model the nonstationarity $p(\mathbf{s}|\mathbf{u})$. We present $\textsc{Mse}[\bar{\mathbf{x}}_{\text{future}}]$ results for prediction steps $T_{\text{pred}}\in\{2,4,8\}$ in Fig. 3, together with the $\textsc{Mcc}[\bar{\mathbf{z}}_{\text{train}}]$ values in the legend. We see that using $K$ different Mlps and a conditional normalizing flow have similar performance improvements for disentanglement. More importantly, the results suggest that as we predict for longer horizons (e.g. 8 steps), a higher Mcc on the latent states lead to a larger improvement on the future prediction performance. Finally, in Fig. 2, we visualize our model predictions in the data space, highlighting the calibrated uncertainty handling of our probabilistic approach.

Figure: Figure 4: A comparison of MSEs achieved by our approach and two-stage LEAP training.
Refer to caption: x5.png

#### Leap with offline transition training vs. our approach.

Next, we demonstrate the benefits of our learning procedure, which identifies the transition and generative functions jointly.
In principle, one can use an off-the-shelf identifiable representation learning method, such as Leap-Np, identify the ground-truth factors, and fit a state transition function on the inferred latent codes.
Such a transition function could approximate the unknown transition function equally well.
We implement and test this two-stage procedure on our synthetic dataset.
To learn the unknown transitions, we train a three-layer MLP with leaky relu activations on the latent state sequence $\bar{\mathbf{z}}_{\text{train}}$ inferred by Leap-Np (different hyperparameters yielded similar results).
We unroll the learned transition function for an additional $T_{\text{future}}$ time steps, decode all the latent states, and compute the Mse in the data space.
We repeat this three times and compare the average error with the error of our approach in Figure 4.
Our approach achieves smaller errors at all time points, and the gap widens as we predict for longer horizons.
This indicates that identifying the generative and transition functions jointly results in a better model.

### 4.2 Modified Cartpole experiment

Next, we experiment on challenging video data of a physical cartpole system with nonlinear dynamics [Huang et al., , Yao et al., ]. Similarly to the synthetic experiment, we show our method yields more accurate latent system identification and future predictions than baselines (Q1,Q4).

#### Cartpole dataset.

We use the modified cartpole setting in [Huang et al., , Yao et al., ], an adaptation of the environment in OpenAI Gym [Brockman et al., ]. The environment simulates a physical system with a pole attached to a cart from a pivot point. The system has 4 states: position and velocity of the cart, and the angle and the angular velocity of the pole ($[x_{\text{cart}},\theta_{\text{pole}},\dot{x}_{\text{cart}},\dot{\theta}_{
\text{pole}}]$). Observations consist of video sequences of the system and we use the true states only to compute Mcc.
The setup has 5 source domains used for training with the different levels of gravity $g=\{5,10,20,30,40\}$ and mass $m=1.0$.
We use a target domain with $g=90$ to check the methods’ ability to extrapolate.
We observe the system under random binary actions, which move the cart to left or right.
We generate sequences of length $T=T_{0}+T_{\text{dyn}}+T_{\text{future}}=1+8+8=17$, the first frame for the initial state, the next 8 frames for training the model, and the last 8 frames for future prediction. For each source domain, we have 900/100 sequences as train/validation data. For the target domain, we have data sets with different numbers of samples $N_{\text{target}}=\{20,50,1000\}$.

#### Main results.

Table 2 presents accuracies of reconstruction and prediction of the latent states. (Here also predictions are compared in the latent space since the true future states contain pole angles unseen during training, resulting in blurry reconstructions.)
As expected, methods that take complex temporal dependencies into account (Leap-Lin, Leap-Np and Ours) perform better than other baselines. Our model recovers latent states the best, i.e., the inferred states are more strongly correlated with the ground-truth.
Also, in line with the synthetic experiment, more accurate latent states on the initial training steps also lead to higher future prediction accuracy.

#### Adaptation experiment.

Next, we take the converged checkpoints of Leap-Lin and our approach and adapt them to the target domain described earlier.
Because both models claim to identify the latent states consisting of positions and velocities, we expect them to transfer well to a new dataset with the same generative mechanism.
We consider a domain shift ($g=40\rightarrow 90$) and observe that the Mccvalues for both models drop to the range $[0.60-0.65]$.
To mitigate this, we train the models on a small “adaptation dataset” with $N=\{20,50,1000\}$ sequences.
We try adapting (i) only the dynamics (including process noise flow) and (ii) the entire network.
Fig. 6 shows the findings.
First, despite decreasing the loss, Leap-Lin Mccscore never improves irrespective of the dataset size.
Importantly, our model improves its Mccby about $25\%$ with only $50$ adaptation trajectories.
Interestingly, optimizing only the dynamics results in better MCC than optimizing all modules.

**Table 2: Results for the cartpole experiment. N/A: the method cannot predict future.**
|  | Models |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Metrics | Pcl | SlowVae | KalmanVae | Leap-Lin | Leap-Np | Ours |
| $\textsc{Mcc}[\bar{\mathbf{z}}_{\text{train}}]$ | $0.62$ | $0.54$ | $0.60$ | $0.83$ | $0.80$ | $0.95$ |
| $\textsc{Mcc}[\bar{\mathbf{z}}_{\text{future}[8]}]$ | N/A | N/A | $0.55$ | $0.71$ | N/A | $0.93$ |

## 5 Discussion

We have presented the first latent dynamical system that allows for the identification of the unknown transition function, and theoretically proved its identifiability, based on standard assumptions. We evaluated our approach on synthetic data and the Cartpole environment and showed that (i) the estimated latent states correlated strongly with the ground truth, (ii) our method had the highest future prediction accuracy with calibrated uncertainties, and (iii) it could adapt to new environments using a handful of data. The main limitation stems from the identifiability assumptions adopted. For future work, it would be intriguing to examine their violation might influence the final model’s performance. Further, our most realistic empirical demonstration is based on data from a physics simulator, i.e., we leave studies on real-world datasets for future. Finally, demonstrating improved downstream performance from our method, e.g. in model-based policy learning, would be of interest.

## Appendix A Identifiability Theory

In this section, we discuss the identifiability of the latent states and the transition function, and provide the detailed proofs.

We assume a latent dynamical system which is viewed as high-dimensional sensory observations $\mathbf{x}_{1:T}$, where $t$ is the time index and $\mathbf{x}_{t}\in\mathbb{R}^{D}$. We assume a sequence of latent states $\mathbf{z}_{1:T}$, with $\mathbf{z}_{t}\in\mathbb{R}^{K}$, are instantaneously mapped to observations via a generative function ${\mathbf{g}}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{D}$:

$$ $\mathbf{x}_{t}={\mathbf{g}}(\mathbf{z}_{t}).$ (13) $$

The latent states $\mathbf{z}_{1:T}$ evolve according to Markovian dynamics:

$$ $\mathbf{z}_{t}=\mathbf{f}(\mathbf{z}_{t-1},\mathbf{s}_{t}),$ (14) $$

where $\mathbf{f}:\mathbb{R}^{2K}\rightarrow\mathbb{R}^{K}$ is an auto-regressive transition function and $\mathbf{s}_{t}\in\mathbb{R}^{K}$ corresponds to process noise.

Our aim is to jointly identify the latent states $\mathbf{z}_{1:T}$, the dynamics function $\mathbf{f}$, and the process noise $\mathbf{s}_{1:T}$. We remind that previous works [Klindt et al., , Yao et al., , , Song et al., ] have concentrated on identifying the latent states $\mathbf{z}_{1:T}$, possibly with linear transition approximations, but not a general transition function $\mathbf{f}$.
Yet, without a general $\mathbf{f}$, the methods can estimate the underlying states only when corresponding observations are provided or provide simplistic approximations in their absence. Hence, they cannot predict complex future behavior reliably.

Notice that learning a provably identifiable transition function $\mathbf{f}:\mathbb{R}^{2K}\rightarrow\mathbb{R}^{K}$ is not straightforward, since the transition function is not injective. A naive solution can be to simply use a plug-in method [Yao et al., , ] for identifying the latents and then fitting a transition function $\mathbf{f}$ on the estimated latents, however, we show empirically in our experiments that it leads to poor prediction accuracy for the future behavior.

### A.1 Nonlinear Ica

The nonlinear Ica assumes that the data is generated from independent latent variables $\mathbf{z}$ with a nonlinear generative function ${\mathbf{g}}$, following Eq. 13. It is well-known to be non-identifiable for i.i.d. data [Hyvärinen and Pajunen, , Locatello et al., ]. Recent seminal works [Hyvarinen and Morioka, , , Hyvarinen et al., ] showed that autocorrelation and nonstationarity existent in non-i.i.d. data can be exploited to identify latent variables in an unsupervised way.
Compared to the vanilla Ica that considers independence only along latent dimensions, the idea of these works is to introduce additional independence constraints reflecting the existent structure in the data. These additional constraints are formulated mathematically as identifiability assumptions, which restrict the space of the generative function ${\mathbf{g}}$ and the space of the latent prior $p_{\mathbf{z}}$ [Hyvärinen et al., , Xi and Bloem-Reddy, ]. The key insight is that, after sufficiently constraining the latent prior $p_{\mathbf{z}}$ using such assumptions, identifying the latent variables $\mathbf{z}_{t}$ and identifying the injective generative function ${\mathbf{g}}$ become equivalent tasks [Xi and Bloem-Reddy, ].

### A.2 Augmented dynamics for identifiable systems

To identify the transition function $\mathbf{f}$ such that $\mathbf{z}_{t}=\mathbf{f}(\mathbf{z}_{t-1},\mathbf{s}_{t})$, we will use the same insight: After sufficiently constraining the noise prior $p_{\mathbf{s}}$; given an identifiable latent pair $(\mathbf{z}_{t-1},\mathbf{z}_{t})$, identifying the noise variables $\mathbf{s}_{t}$ and identifying the bijective dynamics function $\mathbf{f}$ should be equivalent. Hence, in addition to the identifiability assumptions restricting the space of the generative function ${\mathbf{g}}$ and the space of the latent prior $p_{\mathbf{z}}$, we will further restrict the space of the dynamics function $\mathbf{f}$, and the space of the noise prior $p_{\mathbf{s}}$.

First, let us note that the identifiability of the process noise $\mathbf{s}_{t}$ is not trivial since the dynamics function $\mathbf{f}:\mathbb{R}^{2K}\rightarrow\mathbb{R}^{K}$ is not an injective function and hence it does not have an inverse. Following the independent innovation analysis (Iia) framework Morioka et al. [], we trivially augment the image space of the transition function and denote the bijective augmented function by ${\mathbf{f}_{\texttt{aug}}}:\mathbb{R}^{2K}\rightarrow\mathbb{R}^{2K}$:

$$ $\displaystyle\begin{bmatrix}\mathbf{z}_{t}\\ \mathbf{z}_{t-1}\end{bmatrix}={\mathbf{f}_{\texttt{aug}}}\left(\begin{bmatrix} \mathbf{s}_{t}\\ \mathbf{z}_{t-1}\end{bmatrix}\right)=\begin{bmatrix}\mathbf{f}(\mathbf{z}_{t-1 },\mathbf{s}_{t})\\ \mathbf{z}_{t-1}\end{bmatrix}$ (15) $$

Contrary to Morioka et al. [], which use an augmented autoregressive model on observations $(\mathbf{x}_{t-1},\mathbf{x}_{t})$, our formulation captures the functional dependence between a latent pair $(\mathbf{z}_{t-1},\mathbf{z}_{t})$ and the process noise $\mathbf{s}_{t}$.

Next, we make the standard assumption in the temporal identifiability literature [Klindt et al., , Yao et al., , , Song et al., ] that each dimension of the transition function $\{f_{k}\}_{k=1}^{K}$ is influenced by a single process noise variable $s_{kt}$. The output is a single latent variable $z_{kt}$:

$$ $\displaystyle z_{kt}$ $\displaystyle=f_{k}(\mathbf{z}_{t-1},s_{kt}),\quad\text{for }k\in 1,\ldots,K \text{ and }t\in 1,\ldots,T.$ (16) $$

Notice that this does not impose a limitation on the generative model, it just creates a segmentation between noise variables and latent variables. For example, if this assumption is violated and there exists a noise variable $s_{kt}$ that affects both $z_{it}$ and $z_{jt}$ with $i\neq j$, then the noise variable $s_{kt}$ can instead be modeled as a latent variable $z_{kt}$.

We re-state the full generative model for completeness:

$$ $\displaystyle\mathbf{z}_{0}$ $\displaystyle\sim p_{\mathbf{z}_{0}}(\mathbf{z}_{0}),$ # initial state (17) $\displaystyle\mathbf{s}_{t}$ $\displaystyle\sim p_{\mathbf{s}|\mathbf{u}}(\mathbf{s}_{t}|\mathbf{u})=\prod_{ k}p_{s_{k}|\mathbf{u}}(s_{kt}|\mathbf{u}),\quad\forall t\in 1,\ldots,T,$ # process noise (18) $\displaystyle\begin{bmatrix}\mathbf{z}_{t}\\ \mathbf{z}_{t-1}\end{bmatrix}$ $\displaystyle={\mathbf{f}_{\texttt{aug}}}\left(\begin{bmatrix}\mathbf{s}_{t}\\ \mathbf{z}_{t-1}\end{bmatrix}\right)=\begin{bmatrix}\mathbf{f}(\mathbf{z}_{t-1 },\mathbf{s}_{t})\\ \mathbf{z}_{t-1}\end{bmatrix},\quad\forall t\in 1,\ldots,T,$ # state transition (19) $\displaystyle\begin{bmatrix}\mathbf{x}_{t}\\ \mathbf{x}_{t-1}\end{bmatrix}$ $\displaystyle={{\mathbf{g}}_{\texttt{aug}}}\left(\begin{bmatrix}\mathbf{z}_{t} \\ \mathbf{z}_{t-1}\end{bmatrix}\right)=\begin{bmatrix}{\mathbf{g}}(\mathbf{z}_{t })\\ {\mathbf{g}}(\mathbf{z}_{t-1})\end{bmatrix},\qquad\forall t\in 2,\ldots,T.$ # observation mapping (20) $$

where $\mathbf{u}$ is an auxiliary variable, which modulates the noise distribution $p_{\mathbf{s}|\mathbf{u}}$.

### A.3 Proof of Theorem 1: Identifiability of the latent states 𝐳 t subscript 𝐳 𝑡 \mathbf{z}_{t} bold_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT

This result is already shown in [Yao et al., , Appendix A.3.2]. Here, we follow Klindt et al. [], Yao et al. [, ] and repeat their results in our notation as we also make use of this result in Section A.4.

The injective functions ${\mathbf{g}},{\hat{{\mathbf{g}}}}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{D}$ are bijective between the latent space $\mathbb{R}^{K}$ and the observation space $\mathcal{X}\subset\mathbb{R}^{D}$. We denote the inverse functions from the restricted observation space to the latent space by ${\mathbf{g}}^{-1},{\hat{{\mathbf{g}}}}^{-1}$. This is also implicitly assumed in [Klindt et al., , Yao et al., , , Song et al., ].

We have

$$ $\displaystyle\mathbf{x}_{t}={\hat{{\mathbf{g}}}}(\hat{\mathbf{z}}_{t})=\bigg{( }({\mathbf{g}}\circ\underbrace{{\mathbf{g}}^{-1})\circ{\hat{{\mathbf{g}}}}}_{{ \mathbf{h}}}\bigg{)}(\hat{\mathbf{z}}_{t})\implies{\hat{{\mathbf{g}}}}={ \mathbf{g}}\circ{\mathbf{h}}\implies\mathbf{z}_{t}={\mathbf{h}}(\hat{\mathbf{z }}_{t}).$ (21) $$

The function ${\mathbf{h}}:\hat{\mathbf{z}}_{t}\mapsto\mathbf{z}_{t}$ maps the learned latents to the ground-truth latents. To show it is bijective, we need to show it is both injective and surjective. Following Klindt et al. [], it is injective since it is a composition of injective functions. Assume it is not surjective, then there exists a neighborhood $\mathbf{U_{z}}$ for which ${\mathbf{g}}(\mathbf{U_{z}})\notin{\hat{{\mathbf{g}}}}(\mathbb{R}^{K})$. This implies that the neighborhood of images generated by ${\mathbf{g}}(\mathbf{U_{z}})$ has zero density under the learned observation density $p_{{\hat{{\mathbf{g}}}_{\texttt{aug}}},{\hat{\mathbf{f}}_{\texttt{aug}}},\hat{
p}_{\mathbf{s}|\mathbf{u}}}({\mathbf{g}}(\mathbf{U_{z}}))=0$, while having non-zero density under the ground-truth observation density $p_{{\mathbf{g}},{\mathbf{f}_{\texttt{aug}}}}(\mathbf{x})$: $p_{{{\mathbf{g}}_{\texttt{aug}}},{\mathbf{f}_{\texttt{aug}}},p_{\mathbf{s}|
\mathbf{u}}}({\mathbf{g}}(\mathbf{U_{z}}))>0$. This contradicts the assumption that the observation densities match everywhere. Then, ${\mathbf{h}}$ is surjective.

We perform change of variables on the conditional latent density:

$$ $\displaystyle\log p(\hat{\mathbf{z}}_{t}|\hat{\mathbf{z}}_{t-1},\mathbf{u})$ $\displaystyle=\log p(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{u})+\log|{\mathbf {H}}_{t}|,$ (22) $\displaystyle\sum_{k=1}^{K}\underbrace{\log p(\hat{z}_{kt}|\hat{\mathbf{z}}_{t -1},\mathbf{u})}_{\hat{\eta}_{k}(\hat{z}_{kt},\mathbf{u})}$ $\displaystyle=\sum_{k=1}^{K}\underbrace{\log p(z_{kt}|\mathbf{z}_{t-1},\mathbf {u})}_{\eta_{k}(z_{kt},\mathbf{u})}+\log|{\mathbf{H}}_{t}|$ (23) $\displaystyle\sum_{k=1}^{K}\hat{\eta}_{k}(\hat{z}_{kt},\mathbf{u})$ $\displaystyle=\sum_{k=1}^{K}\eta_{k}(z_{kt},\mathbf{u})+\log|{\mathbf{H}}_{t}|$ (24) $$

where ${\mathbf{H}}_{t}=\mathbf{J_{h}}(\hat{\mathbf{z}}_{t})$ is the Jacobian matrix of ${\mathbf{h}}$ evaluated at $\hat{\mathbf{z}}_{t}$. First, we take derivatives of both sides with respect to $\hat{z}_{it}$:

$$ $\displaystyle\hat{\eta}_{i}(\hat{z}_{it},\mathbf{u})$ $\displaystyle=\sum_{k=1}^{K}\frac{\partial\eta_{k}(z_{kt},\mathbf{u})}{ \partial z_{kt}}\frac{\partial z_{kt}}{\partial\hat{z}_{it}}+\frac{\partial \log|{\mathbf{H}}_{t}|}{\partial\hat{z}_{it}}$ (25) $$

Second, take derivatives with respect to $\hat{z}_{jt}$:

$$ $\displaystyle 0$ $\displaystyle=\sum_{k=1}^{K}\left(\frac{\partial^{2}\eta_{k}(z_{kt},\mathbf{u} )}{\partial z^{2}_{kt}}\frac{\partial z_{kt}}{\partial\hat{z}_{it}}\frac{ \partial z_{kt}}{\partial\hat{z}_{jt}}+\frac{\partial\eta_{k}(z_{kt},\mathbf{u })}{\partial z_{kt}}\frac{\partial z^{2}_{kt}}{\partial\hat{z}_{it}\partial \hat{z}_{jt}}\right)+\frac{\partial\log|{\mathbf{H}}_{t}|}{\partial\hat{z}_{it }\partial\hat{z}_{jt}}$ (26) $$

Lastly, take derivatives with respect to $u_{l}$:

$$ $\displaystyle 0$ $\displaystyle=\sum_{k=1}^{K}\left(\frac{\partial^{3}\eta_{k}(z_{kt},\mathbf{u} )}{\partial z^{2}_{kt}\partial u_{l}}\frac{\partial z_{kt}}{\partial\hat{z}_{ it}}\frac{\partial z_{kt}}{\partial\hat{z}_{jt}}+\frac{\partial^{2}\eta_{k}(z_ {kt},\mathbf{u})}{\partial z_{kt}\partial u_{l}}\frac{\partial z^{2}_{kt}}{ \partial\hat{z}_{it}\partial\hat{z}_{jt}}\right),$ (27) $\displaystyle=\sum_{k=1}^{K}\left(\frac{\partial^{3}\eta_{k}(z_{kt},\mathbf{u} )}{\partial z^{2}_{kt}\partial u_{l}}[{\mathbf{H}}_{t}]_{ki}[{\mathbf{H}}_{t}] _{kj}+\frac{\partial^{2}\eta_{k}(z_{kt},\mathbf{u})}{\partial z_{kt}\partial u _{l}}\frac{\partial z^{2}_{kt}}{\partial\hat{z}_{it}\partial\hat{z}_{jt}} \right),$ (28) $$

since the Jacobian ${\mathbf{H}}_{t}$ does not depend on $\mathbf{u}$. Using the sufficient variability assumption (A4) for the latent states $\mathbf{z}_{t}$, we can plug in $2K$ values of $\mathbf{u}_{1},\ldots,\mathbf{u}_{2K}$ for which the partial derivatives of the log conditional density $\eta_{k}(z_{kt},\mathbf{u})$ form linearly independent vectors $\mathbf{v}(\mathbf{z}_{t},\mathbf{u})$. We see that the coefficients of these linearly independent vectors have to be zero: $[{\mathbf{H}}_{t}]_{ki}[{\mathbf{H}}_{t}]_{kj}=0$. This implies that the Jacobian matrix ${\mathbf{H}}_{t}$ of the transformation $\mathbf{z}_{t}={\mathbf{h}}(\hat{\mathbf{z}}_{t})$ has at most $1$ nonzero element in its rows. Therefore, the learned latents $\hat{\mathbf{z}}_{t}$ are equivalent to the ground-truth latents $\mathbf{z}_{t}$ up to permutations and invertible, element-wise nonlinear transformations.

### A.4 Proof of Theorem 2: Identifiability of the latent transition 𝐟 𝐟 \mathbf{f} bold_f

Here, we prove our main theoretical contribution, Theorem 2. We start by writing the generative process for the observation pair $(\mathbf{x}_{t-1},\mathbf{x}_{t})$:

$$ $\displaystyle\begin{bmatrix}\mathbf{x}_{t}\\ \mathbf{x}_{t-1}\end{bmatrix}$ $\displaystyle=({\hat{{\mathbf{g}}}_{\texttt{aug}}}\circ{\hat{\mathbf{f}}_{ \texttt{aug}}})\left(\begin{bmatrix}\hat{\mathbf{s}}_{t}\\ \hat{\mathbf{z}}_{t-1}\end{bmatrix}\right)=\begin{bmatrix}({\hat{{\mathbf{g}}} }\circ{\hat{\mathbf{f}}})(\hat{\mathbf{z}}_{t-1},\hat{\mathbf{s}}_{t})\\ {\hat{{\mathbf{g}}}}(\hat{\mathbf{z}}_{t-1})\end{bmatrix}$ (29) $\displaystyle=({{\mathbf{g}}_{\texttt{aug}}}\circ{\mathbf{f}_{\texttt{aug}}}) \circ\underbrace{({{\mathbf{g}}_{\texttt{aug}}}\circ{\mathbf{f}_{\texttt{aug}} })^{-1}\circ({\hat{{\mathbf{g}}}_{\texttt{aug}}}\circ{\hat{\mathbf{f}}_{ \texttt{aug}}})}_{{{\mathbf{k}}_{\texttt{aug}}}}\left(\begin{bmatrix}\hat{ \mathbf{s}}_{t}\\ \hat{\mathbf{z}}_{t-1}\end{bmatrix}\right)$ (30) $$

The function ${{\mathbf{k}}_{\texttt{aug}}}:\mathbb{R}^{2K}\rightarrow\mathbb{R}^{2K}$ maps the learned pair $(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1})$ to the ground-truth pair $(\mathbf{s}_{t},\mathbf{z}_{t-1})$:

$$ $\displaystyle\begin{bmatrix}\mathbf{s}_{t}\\ \mathbf{z}_{t-1}\end{bmatrix}$ $\displaystyle={{\mathbf{k}}_{\texttt{aug}}}\left(\begin{bmatrix}\hat{\mathbf{s }}_{t}\\ \hat{\mathbf{z}}_{t-1}\end{bmatrix}\right)=\begin{bmatrix}{\mathbf{k}}_{1}( \hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1})\\ {\mathbf{k}}_{2}(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1})\end{bmatrix}.$ (31) $$

Similar to the function ${\mathbf{h}}$ being bijective, it follows that ${{\mathbf{k}}_{\texttt{aug}}}$ is also bijective.

Now, we want to show that the augmented function ${{\mathbf{k}}_{\texttt{aug}}}$ decomposes into invertible block-wise functions ${\mathbf{k}}_{1},{\mathbf{k}}_{2}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{K}$ such that (i) ${\mathbf{k}}_{1}$ only depends on $\hat{\mathbf{s}}_{t}$, i.e., $\mathbf{s}_{t}={\mathbf{k}}_{1}(\hat{\mathbf{s}}_{t})$ and (ii) ${\mathbf{k}}_{2}$ only depends on $\hat{\mathbf{z}}_{t-1}$, i.e., $\mathbf{z}_{t-1}={\mathbf{k}}_{2}(\hat{\mathbf{z}}_{t-1})$. It is easy to show (ii), since ${\mathbf{k}}_{2}=\mathrm{id}_{\mathbf{z}}\circ{\mathbf{g}}^{-1}\circ{\hat{{
\mathbf{g}}}}\circ\mathrm{id}_{\hat{\mathbf{z}}}={\mathbf{h}}$ and we have already shown in Section A.3 that the function ${\mathbf{h}}:\hat{\mathbf{z}}_{t}\mapsto\mathbf{z}_{t}$ is equal to ${\mathbf{h}}={\mathbf{g}}^{-1}\circ{\hat{{\mathbf{g}}}}$ and bijective.

$$ $\displaystyle\begin{bmatrix}\mathbf{s}_{t}\\ \mathbf{z}_{t-1}\end{bmatrix}$ $\displaystyle={{\mathbf{k}}_{\texttt{aug}}}\left(\begin{bmatrix}\hat{\mathbf{s }}_{t}\\ \hat{\mathbf{z}}_{t-1}\end{bmatrix}\right)=\begin{bmatrix}{\mathbf{k}}_{1}( \hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1})\\ {\mathbf{h}}(\hat{\mathbf{z}}_{t-1})\end{bmatrix}.$ (32) $$

To show (i), we start by performing change of variables for the transformation ${{\mathbf{k}}_{\texttt{aug}}}:(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1})
\mapsto(\mathbf{s}_{t},\mathbf{z}_{t-1})$:

$$ $\displaystyle\log p(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1}|\mathbf{u})$ $\displaystyle=\log p({{\mathbf{k}}_{\texttt{aug}}}(\hat{\mathbf{s}}_{t},\hat{ \mathbf{z}}_{t-1})|\mathbf{u})+\log|\mathbf{J_{{{\mathbf{k}}_{\texttt{aug}}}}} (\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1})|,$ (33) $\displaystyle=\log p([{\mathbf{k}}_{1}(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{ t-1}),{\mathbf{h}}(\hat{\mathbf{z}}_{t-1})]|\mathbf{u})+\log|\mathbf{J_{{{ \mathbf{k}}_{\texttt{aug}}}}}(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1})|,$ (34) $$

where $\mathbf{J_{{{\mathbf{k}}_{\texttt{aug}}}}}(\hat{\mathbf{s}}_{t},\hat{\mathbf{z
}}_{t-1})$ is the Jacobian matrix for the augmented function ${{\mathbf{k}}_{\texttt{aug}}}$ evaluated at $(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1})$. As the process noise is temporally independent given $\mathbf{u}$, $\hat{\mathbf{s}}_{t}\perp\!\!\!\perp\hat{\mathbf{z}}_{t-1}|\mathbf{u}$ and $\mathbf{s}_{t}\perp\!\!\!\perp\mathbf{z}_{t-1}|\mathbf{u}$, we factorize the densities in Eq. 34:

$$ $\displaystyle\log p(\hat{\mathbf{s}}_{t}|\mathbf{u})+\log p(\hat{\mathbf{z}}_{ t-1}|\mathbf{u})$ $\displaystyle=\log p({\mathbf{k}}_{1}(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t -1})|\mathbf{u})+\log p({\mathbf{h}}(\hat{\mathbf{z}}_{t-1})|\mathbf{u})+\log| \mathbf{J_{{{\mathbf{k}}_{\texttt{aug}}}}}(\hat{\mathbf{s}}_{t},\hat{\mathbf{z }}_{t-1})|.$ (35) $$

The Jacobian $\mathbf{J_{{{\mathbf{k}}_{\texttt{aug}}}}}$ is upper block-diagonal since $\mathbf{z}_{t-1}$ does not depend on $\hat{\mathbf{s}}_{t}$: $\mathbf{J_{{{\mathbf{k}}_{\texttt{aug}}}}}=\begin{bmatrix}\frac{\partial
\mathbf{s}_{t}}{\partial\hat{\mathbf{s}}_{t}}&\mathbf{*}\\
\mathbf{0}&{\mathbf{H}}_{t}\end{bmatrix}$ and its log determinant factorizes $\log|\mathbf{J_{{{\mathbf{k}}_{\texttt{aug}}}}}(\hat{\mathbf{s}}_{t},\hat{
\mathbf{z}}_{t-1})|=\log|{\mathbf{H}}_{t}|+\log|\frac{\partial\mathbf{s}_{t}}{
\partial\hat{\mathbf{s}}_{t}}|$:

$$ $\displaystyle\log p(\hat{\mathbf{s}}_{t}|\mathbf{u})+\log p(\hat{\mathbf{z}}_{ t-1}|\mathbf{u})$ $\displaystyle=\log p({\mathbf{k}}_{1}(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t -1})|\mathbf{u})+\log p({\mathbf{h}}(\hat{\mathbf{z}}_{t-1})|\mathbf{u})+\log| {\mathbf{H}}_{t}|+\log|\frac{\partial\mathbf{s}_{t}}{\partial\hat{\mathbf{s}}_ {t}}|.$ (36) $$

In addition, the noise is conditionally independent over its dimensions given $\mathbf{u}$. Therefore, we can further factorize the densities $p(\hat{\mathbf{s}}_{t}|\mathbf{u})=\prod_{k}p(\hat{s}_{kt}|\mathbf{u})$ and $p({\mathbf{k}}_{1}(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1})|\mathbf{u})=p(
\mathbf{s}_{t}|\mathbf{u})=\prod_{k}p(s_{kt}|\mathbf{u})$ with $\mathbf{s}_{t}={\mathbf{k}}_{1}(\hat{\mathbf{s}}_{t},\hat{\mathbf{z}}_{t-1})$:

$$ $\displaystyle\sum_{k}\underbrace{\log p(\hat{s}_{kt}|\mathbf{u})}_{\hat{q}_{k} (\hat{s}_{kt},\mathbf{u})}+\log p(\hat{\mathbf{z}}_{t-1}|\mathbf{u})$ $\displaystyle=\sum_{k}\underbrace{\log p(s_{kt}|\mathbf{u})}_{q_{k}(s_{kt}, \mathbf{u})}+\log p({\mathbf{h}}(\hat{\mathbf{z}}_{t-1})|\mathbf{u})+\log|{ \mathbf{H}}_{t}|+\log|\frac{\partial\mathbf{s}_{t}}{\partial\hat{\mathbf{s}}_{ t}}|,$ (37) $$

We take the derivative of both sides with respect to $\hat{s}_{it}$:

$$ $\displaystyle\frac{\partial\hat{q}_{i}(\hat{s}_{it},\mathbf{u})}{\partial\hat{ s}_{it}}$ $\displaystyle=\sum_{k}\frac{\partial q_{k}(s_{kt},\mathbf{u})}{\partial s_{kt} }\frac{\partial s_{kt}}{\partial\hat{s}_{it}}+\frac{\partial\log|\frac{ \partial\mathbf{s}_{t}}{\partial\hat{\mathbf{s}}_{t}}|}{\partial\hat{s}_{it}}.$ (38) $$

Next, we take the derivative with respect to $u_{l}$ with $l$ being an arbitrary dimension:

$$ $\displaystyle\frac{\partial^{2}\hat{q}_{i}(\hat{s}_{it},\mathbf{u})}{\partial \hat{s}_{it}\partial u_{l}}$ $\displaystyle=\sum_{k}\frac{\partial^{2}q_{k}(s_{kt},\mathbf{u})}{\partial s_{ kt}\partial u_{l}}\frac{\partial s_{kt}}{\partial\hat{s}_{it}},$ (39) $$

since $|\frac{\partial\mathbf{s}_{t}}{\partial\hat{\mathbf{s}}_{t}}|$ does not depend on $\mathbf{u}$. Lastly, take the derivative of both sides with respect to $\hat{z}_{j,t-1}$:

$$ $\displaystyle 0$ $\displaystyle=\sum_{k}\left(\frac{\partial^{3}q_{k}(s_{kt},\mathbf{u})}{ \partial s^{2}_{kt}\partial u_{l}}\frac{\partial s_{kt}}{\partial\hat{s}_{it}} \frac{\partial s_{kt}}{\partial\hat{z}_{j,t-1}}+\frac{\partial^{2}q_{k}(s_{kt} ,\mathbf{u})}{\partial s_{kt}\partial u_{l}}\frac{\partial^{2}s_{kt}}{\partial \hat{s}_{it}\partial\hat{z}_{j,t-1}}\right).$ (40) $$

Inspecting the Eq. 40, to ensure the sufficient variability assumption (A5) for the process noise $\mathbf{s}_{t}$, the term $\frac{\partial s_{kt}}{\partial\hat{s}_{it}}\frac{\partial s_{kt}}{\partial
\hat{z}_{j,t-1}}=0$. Following a similar reasoning with Morioka et al. [], this implies that any dimension of $\mathbf{s}_{t}$ does not depend on $\hat{\mathbf{s}}_{t}$ and $\hat{\mathbf{z}}_{t-1}$ at the same time. Since $\mathbf{s}_{t}\perp\!\!\!\perp\mathbf{z}_{t-1}|\mathbf{u}$ and $\mathbf{z}_{t-1}={\mathbf{h}}(\hat{\mathbf{z}}_{t-1})$, $\mathbf{s}_{t}$ has to depend solely on $\hat{\mathbf{s}}_{t}$: $\mathbf{s}_{t}={\mathbf{k}}_{1}(\hat{\mathbf{z}}_{t-1},\hat{\mathbf{s}}_{t})={
\mathbf{k}}(\hat{\mathbf{s}}_{t})$. We conclude that the augmented function ${{\mathbf{k}}_{\texttt{aug}}}$ decomposes into invertible block-wise functions ${\mathbf{k}}$ and ${\mathbf{h}}$: ${{\mathbf{k}}_{\texttt{aug}}}=[{\mathbf{k}},{\mathbf{h}}]$

Now, let’s get back to Eq. 39. Denote the Jacobian matrix of function ${\mathbf{k}}$ by $\mathbf{J_{k}}$ and its evaluation at $\hat{\mathbf{s}}_{t}$ by $\mathbf{J_{k}}(\hat{\mathbf{s}}_{t})={\mathbf{K}}_{t}$. Take the derivative of both sides with respect to $\hat{s}_{mt}$ for some index $m$:

$$ $\displaystyle 0$ $\displaystyle=\sum_{k}\left(\frac{\partial^{3}q_{k}(s_{kt},\mathbf{u})}{ \partial s^{2}_{kt}\partial u_{l}}[{\mathbf{K}}_{t}]_{ki}[{\mathbf{K}}_{t}]_{ km}+\frac{\partial^{2}q_{k}(s_{kt},\mathbf{u})}{\partial s_{kt}\partial u_{l}} \frac{\partial^{2}s_{kt}}{\partial\hat{s}_{it}\partial\hat{s}_{mt}}\right).$ (41) $$

Inspecting the Eq. 41, we see that to ensure the sufficient variability assumption (A5), the product $[{\mathbf{K}}_{t}]_{ki}[{\mathbf{K}}_{t}]_{km}=0$. This implies that each dimension $s_{kt}$ of the true latent state depends only on a single dimension of the learned process noise $\hat{\mathbf{s}}_{t}$. Hence, the function ${\mathbf{k}}$ is equal to a composition of permutation and element-wise, invertible nonlinear transformation: ${\mathbf{k}}=\pi\circ T$.

Following Eq. 30, we can write the relationship between the augmented functions as:

$$ $\displaystyle{{\mathbf{g}}_{\texttt{aug}}}\circ{\mathbf{f}_{\texttt{aug}}} \circ{{\mathbf{k}}_{\texttt{aug}}}$ $\displaystyle={\hat{{\mathbf{g}}}_{\texttt{aug}}}\circ{\hat{\mathbf{f}}_{ \texttt{aug}}},$ (42) $\displaystyle\underbrace{{\hat{{\mathbf{g}}}_{\texttt{aug}}}^{-1}\circ{{ \mathbf{g}}_{\texttt{aug}}}}_{{\mathbf{h}}^{-1}_{\texttt{aug}}}\circ{\mathbf{f }_{\texttt{aug}}}\circ{{\mathbf{k}}_{\texttt{aug}}}$ $\displaystyle={\hat{\mathbf{f}}_{\texttt{aug}}},$ (43) $\displaystyle{\mathbf{h}}^{-1}_{\texttt{aug}}\circ{\mathbf{f}_{\texttt{aug}}} \circ{{\mathbf{k}}_{\texttt{aug}}}$ $\displaystyle={\hat{\mathbf{f}}_{\texttt{aug}}},$ (44) $$

where ${\mathbf{h}}^{-1}_{\texttt{aug}}=[{\mathbf{h}}^{-1},{\mathbf{h}}^{-1}]$. We have shown that both ${\mathbf{h}}$ and ${\mathbf{k}}$ are compositions of permutations and element-wise invertible transformations. Hence, the augmented transition function ${\hat{\mathbf{f}}_{\texttt{aug}}}$ is equal to the true augmented transition function ${\mathbf{f}_{\texttt{aug}}}$ up to compositions of permutations and element-wise transformations.

### A.5 Alternative Versions of Sufficient Variability Assumption

If the variable $\mathbf{u}$ is an observed categorical variable (e.g., domain indicator), the assumptions (A4, A5) can be written in an alternative form without partial derivatives with respect to $u_{l}$, similar to Hyvarinen et al. [], Yao et al. []. For example, for the latent states $\mathbf{z}_{t}$, the alternative version of the (A4) takes the form:

- •
Sufficient variability of latent states for a categorical $\mathbf{u}$ [Yao et al., ]. For any $\mathbf{z}_{t}$, there exist some $2K+1$ values for $\mathbf{u}$: $\mathbf{u}_{1},\ldots,\mathbf{u}_{2K}$, such that the $2K$ vectors $\mathbf{v}(\mathbf{z}_{t},\mathbf{u}_{j+1})-\mathbf{v}(\mathbf{z}_{t},\mathbf{
u}_{j})$ with $j=0,1,\ldots,2K$, are linearly independent where
$\displaystyle\mathbf{v}(\mathbf{z}_{t},\mathbf{u})=\left(\frac{\partial\eta_{1
}(z_{1t},\mathbf{u})}{\partial z_{1t}},\cdots,\frac{\partial\eta_{K}(z_{Kt},
\mathbf{u})}{\partial z_{Kt}},\frac{\partial^{2}\eta_{1}(z_{1t},\mathbf{u})}{
\partial z^{2}_{1t}},\cdots,\frac{\partial^{2}\eta_{K}(z_{Kt},\mathbf{u})}{
\partial z^{2}_{Kt}}\right)\in\mathbb{R}^{2K}.$
(46)

A similar categorical version is provided in [Hyvarinen et al., , Assumption 3], while the continuous version is provided in the same work [Hyvarinen et al., , Appendix D].

## Appendix B Variational inference

Similar to previous works [Yao et al., , ], we want to maximize the marginal log-likelihood $\log p(\mathbf{x}_{1:T}|\mathbf{u})$ that is obtained by marginalizing over the latent states $\mathbf{z}_{1:T}$ and process noise $\mathbf{s}_{1:T}$:

$$ $\displaystyle\log p_{\theta}(\mathbf{x}_{1:T}|\mathbf{u})=\log\int_{\mathbf{z} ,\mathbf{s}}p_{\theta}(\mathbf{x}_{1:T},\mathbf{z}_{0:T},\mathbf{s}_{1:T}| \mathbf{u})\mathop{}\!\mathrm{d}\mathbf{z}_{0:T}\mathop{}\!\mathrm{d}\mathbf{s }_{1:T},$ (47) $$

where we decompose the joint distribution as follows:

$$ $\displaystyle p_{\theta}(\mathbf{x}_{1:T},\mathbf{z}_{0:T},\mathbf{s}_{1:T}| \mathbf{u})$ $\displaystyle=p_{\theta}(\mathbf{z}_{0})\prod_{t=1}^{T}p_{\theta}(\mathbf{s}_{ t}|\mathbf{u})\underbrace{p_{\theta}(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{s }_{t})}_{\delta\left(\mathbf{z}_{t}-\mathbf{f}(\mathbf{s}_{t},\mathbf{z}_{t-1} )\right)}p_{\theta}(\mathbf{x}_{t}|\mathbf{z}_{t}).$ (48) $$

Note that the state transitions $p_{\theta}(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{s}_{t})$ are assumed to be deterministic. The above integral is intractable due to non-linear dynamics $\mathbf{f}$ and observation ${\mathbf{g}}$ functions. As typically done with the deep latent variable models, we approximate the log marginal likelihood by a variational lower bound, i.e., we introduce an amortized approximate posterior distribution $q_{\phi}(\mathbf{z}_{0:T},\mathbf{s}_{1:T}|\mathbf{x}_{1:T},\mathbf{u})$ that decomposes as follows:

$$ $\displaystyle q_{\phi}(\mathbf{z}_{0:T},\mathbf{s}_{1:T}|\mathbf{x}_{1:T}, \mathbf{u})=q_{\phi}(\mathbf{z}_{0}|\mathbf{x}_{1:T},\mathbf{u})\prod_{t=1}^{T }\underbrace{q_{\phi}(\mathbf{s}_{t}|\mathbf{z}_{0:t-1},\mathbf{s}_{1:t-1}, \mathbf{x}_{1:T},\mathbf{u})}_{q_{\phi}(\mathbf{s}_{t}|\mathbf{z}_{t-1}, \mathbf{x}_{1:t})}\underbrace{q_{\phi}(\mathbf{z}_{t}|\mathbf{z}_{0:t-1}, \mathbf{s}_{1:t},\mathbf{u})}_{q_{\phi}(\mathbf{z}_{t}|\mathbf{z}_{t-1}, \mathbf{s}_{t})}$ (49) $$

We simplify the variational posterior $q(\mathbf{s}_{t}|\cdot)$ as $q_{\phi}(\mathbf{s}_{t}|\mathbf{z}_{0:t-1},\mathbf{s}_{1:t-1},\mathbf{x}_{1:T}
,\mathbf{u})=q_{\phi}(\mathbf{s}_{t}|\mathbf{z}_{t-1},\mathbf{x}_{1:t})$, corresponding to a filtering distribution. As in the generative model, we choose $q_{\phi}(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{s}_{t})=p_{\theta}(\mathbf{z}
_{t}|\mathbf{z}_{t-1},\mathbf{s}_{t})=\delta\left(\mathbf{z}_{t}-\mathbf{f}(
\mathbf{s}_{t},\mathbf{z}_{t-1})\right)$:

$$ $\displaystyle q_{\phi}(\mathbf{z}_{0:T},\mathbf{s}_{1:T}|\mathbf{x}_{1:T}, \mathbf{u})=q_{\phi}(\mathbf{z}_{0}|\mathbf{x}_{1:T})\prod_{t=1}^{T}q_{\phi}( \mathbf{s}_{t}|\mathbf{z}_{t-1},\mathbf{x}_{1:t})p_{\theta}(\mathbf{z}_{t}| \mathbf{z}_{t-1},\mathbf{s}_{t}),$ (50) $$

where the functional forms of the densities $q_{\phi}(\mathbf{z}_{0}|\cdot)$ and $q_{\phi}(\mathbf{s}_{t}|\cdot)$ are chosen as diagonal Gaussian distributions whose parameters are computed by recurrent neural networks. The variational lower bound takes the following form:

$$ $\displaystyle\mathcal{L}(\theta,\phi)$ $\displaystyle=\operatorname{\mathbb{E}}_{q_{\phi}(\mathbf{z}_{0:T},\mathbf{s}_ {1:T})}\left[\log p_{\theta}(\mathbf{x}_{1:T}|\mathbf{z}_{1:T})+\log\frac{p_{ \theta}(\mathbf{z}_{0:T},\mathbf{s}_{1:T}|\mathbf{u})}{q_{\phi}(\mathbf{z}_{0: T},\mathbf{s}_{1:T}|\mathbf{x}_{1:T},\mathbf{u})}\right]$ (51) $\displaystyle=\underbrace{\sum_{t=1}^{T}\operatorname{\mathbb{E}}_{q_{\phi}( \mathbf{z}_{t})}[\log p_{\theta}(\mathbf{x}_{t}|\mathbf{z}_{t})]}_{\text{ Reconstruction term, }\mathcal{L}_{R}}+\underbrace{\operatorname{\mathbb{E}}_{ q_{\phi}(\mathbf{z}_{0:T},\mathbf{s}_{1:T})}\left[\log\frac{p_{\theta}(\mathbf {z}_{0:T},\mathbf{s}_{1:T}|\mathbf{u})}{q_{\phi}(\mathbf{z}_{0:T},\mathbf{s}_{ 1:T}|\mathbf{x}_{1:T},\mathbf{u})}\right]}_{\text{KL term, }\mathcal{L}_{KL}}.$ (52) $$

The reconstruction term can easily be computed in a variational auto-encoder framework. Below, we provide the derivation of of the KL term:

$$ $\displaystyle\mathcal{L}_{KL}$ $\displaystyle=\operatorname{\mathbb{E}}_{q_{\phi}(\mathbf{z}_{0:T},\mathbf{s}_ {1:T})}\left[\log\frac{p_{\theta}(\mathbf{z}_{0})}{q_{\phi}(\mathbf{z}_{0}| \mathbf{x}_{1:T},\mathbf{u})}+\log\frac{p_{\theta}(\mathbf{z}_{1:T},\mathbf{s} _{1:T}|\mathbf{u})}{q_{\phi}(\mathbf{z}_{1:T},\mathbf{s}_{1:T}|\mathbf{x}_{1:T },\mathbf{u})}\right]$ (53) $\displaystyle=-D_{KL}(q_{\phi}(\mathbf{z}_{0}|\mathbf{x}_{1:T},\mathbf{u})||p_ {\theta}(\mathbf{z}_{0}))+\operatorname{\mathbb{E}}_{q_{\phi}(\mathbf{z}_{0:T} ,\mathbf{s}_{1:T})}\left[\log\frac{p_{\theta}(\mathbf{z}_{1:T},\mathbf{s}_{1:T }|\mathbf{u})}{q_{\phi}(\mathbf{z}_{1:T},\mathbf{s}_{1:T}|\mathbf{x}_{1:T}, \mathbf{u})}\right]$ (54) $\displaystyle=-D_{KL}(q_{\phi}(\mathbf{z}_{0}|\mathbf{x}_{1:T},\mathbf{u})||p_ {\theta}(\mathbf{z}_{0}))+\operatorname{\mathbb{E}}_{q_{\phi}(\mathbf{z}_{0:T} ,\mathbf{s}_{1:T})}\left[\sum_{t=1}^{T}\log\frac{p_{\theta}(\mathbf{s}_{t}| \mathbf{u})p_{\theta}(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{s}_{t})}{q_{\phi }(\mathbf{s}_{t}|\mathbf{z}_{t-1},\mathbf{x}_{1:T},\mathbf{u})q_{\phi}(\mathbf {z}_{t}|\mathbf{z}_{t-1},\mathbf{s}_{t})}\right]$ (55) $\displaystyle=-D_{KL}(q_{\phi}(\mathbf{z}_{0}|\mathbf{x}_{1:T},\mathbf{u})||p_ {\theta}(\mathbf{z}_{0}))+\sum_{t=1}^{T}\operatorname{\mathbb{E}}_{q_{\phi}( \mathbf{z}_{0:T},\mathbf{s}_{1:T})}\left[\log\frac{p_{\theta}(\mathbf{s}_{t}| \mathbf{u})}{q_{\phi}(\mathbf{s}_{t}|\mathbf{z}_{t-1},\mathbf{x}_{1:T},\mathbf {u})}\right]$ (56) $\displaystyle=-D_{KL}(q_{\phi}(\mathbf{z}_{0}|\mathbf{x}_{1:T},\mathbf{u})||p_ {\theta}(\mathbf{z}_{0}))+\sum_{t=1}^{T}\underbrace{\operatorname{\mathbb{E}}_ {q_{\phi}(\mathbf{s}_{t},\mathbf{z}_{t-1}|\mathbf{z}_{t-2},\mathbf{s}_{t-1}, \mathbf{x}_{1:T},\mathbf{u})}\left[\log\frac{p_{\theta}(\mathbf{s}_{t}|\mathbf {u})}{q_{\phi}(\mathbf{s}_{t}|\mathbf{z}_{t-1},\mathbf{x}_{1:T},\mathbf{u})} \right]}_{-\operatorname{\mathbb{E}}_{q_{\phi}(\mathbf{z}_{t-1}|\mathbf{z}_{t- 2},\mathbf{s}_{t-1})}[D_{KL}(q_{\phi}(\mathbf{s}_{t}|\mathbf{z}_{t-1},\mathbf{ x}_{1:T},\mathbf{u})||p_{\theta}(\mathbf{s}_{t}|\mathbf{u}))]}$ (57) $$

Figure: Figure 5: Extended version of Figure 2
Refer to caption: extracted/5647868/figs/uncertainty_plot.png

Figure: Figure 6: We trace the $\mathcal{L}$ and Mcc metrics while adapting LEAP and our approach to small Cartpole adaptation datasets. $x$-axis shows the gradient steps ($\times 1000$)
Refer to caption: extracted/5647868/figs/adaptation_LEAP_elbo_mcc.png

Figure: Figure 7: Diagram of the model architecture: In training, the observation $\mathbf{x}$ is passed through the encoder $\mathbf{e}_{\phi}$ to get the representation $\mathbf{r}$. We learn the distribution over the initial latent state $\mathbf{z}_{0}$ conditional on the representations of the first $T_{ic}$ observations (in the diagram, we show $T_{ic}=1$; in our experiments, we use $T_{\text{ic}}=2$). The latent state is decoded by the decoder $\mathbf{d}_{\theta}$ to produce the predicted observation $\hat{\mathbf{x}}$ (which is trained to match the corresponding actual observation). The next value of the latent state is computed by the transition function $\mathbf{f}_{\theta}$, which depends both on the previous state and on the process noise $\mathbf{s}$. In training, the process noise $\mathbf{s}$ is sampled from the variational posterior that depends on the previous state as well as on the representation created by a recurrent neural network (GRU) that has received up to the current observation. In future prediction, the process noise $\mathbf{s}$ is sampled from the prior, which is a learned normalizing flow.

## Appendix C Architecture and optimization details

We optimize our model with Adam optimizer. We chose all hyperparameters for our method, two versions of LEAP and KalmanVAE with cross-validation. In particular, we performed random search as well as Bayesian optimization over learning rate, weight regularization, the number of layers in all MLPs, and latent dimensionality.

### C.1 Synthetic data experiments

- •
Encoder: Below, the output of encoder_base_layer goes into encoder_rnn_layer and s_encoder.
–
(encoder_base_layer (MLP)):
*
Linear(in_features=8, out_features=64, bias=True)
*
LeakyReLU(negative_slope=0.2)
*
Linear(in_features=64, out_features=64, bias=True)
*
LeakyReLU(negative_slope=0.2)
*
Linear(in_features=64, out_features=64, bias=True)
*
LeakyReLU(negative_slope=0.2)
*
Linear(in_features=64, out_features=64, bias=True)
–
(encoder_rnn_layer): GRU(in=64, hidden_dim=64, output_size=64)
–
(ic_encoder):
*
Linear(in_features=256, out_features=64, bias=True)
*
LeakyReLU(negative_slope=0.2)
*
Linear(in_features=64, out_features=64, bias=True)
*
LeakyReLU(negative_slope=0.2)
*
Linear(in_features=64, out_features=16, bias=True)
–
(s_encoder (MLP)):
*
Linear(in_features=80, out_features=64, bias=True)
*
LeakyReLU(negative_slope=0.2),
*
Linear(in_features=64, out_features=64, bias=True)
*
LeakyReLU(negative_slope=0.2),
*
Linear(in_features=64, out_features=16, bias=True)
- •
Decoder (MLP)
–
Linear(in_features=8, out_features=64, bias=True)
–
LeakyReLU(negative_slope=0.2),
–
Linear(in_features=64, out_features=64, bias=True)
–
LeakyReLU(negative_slope=0.2)
–
Linear(in_features=64, out_features=8, bias=True))
- •
Transition function: Here, we consider 8 different MLPs, each of which has the following architecture:
–
Linear(in_features=17, out_features=64, bias=True)
–
LeakyReLU(negative_slope=0.2),
–
Linear(in_features=64, out_features=64, bias=True)
–
LeakyReLU(negative_slope=0.2)
–
Linear(in_features=64, out_features=1, bias=True))

### C.2 Cartpole experiments

- •
Encoder
–
Conv2d(in_channels=3, num_filter=32, kernel=3, stride=2, pad=1)
–
GeLU()
–
Conv2d(in_channels=32, num_filter=32, kernel=3, stride=2, pad=1)
–
GeLU()
–
Conv2d(in_channels=32, num_filter=64, kernel=3, stride=2, pad=1)
–
GeLU()
–
Conv2d(in_channels=64, num_filter=64, kernel=3, stride=2, pad=1)
–
GeLU()
–
Flatten()
–
Linear(1024,8)
- •
Decoder
–
Linear(8,1024)
–
UnFlatten(4x4x64)
–
ConvTranspose2d(in_channels=64, num_filter=64, kernel=3, stride=2, pad=1, output_padding=1)
–
GeLU()
–
ConvTranspose2d(in_channels=64, num_filter=32, kernel=3, stride=2, pad=1, output_padding=1)
–
GeLU()
–
ConvTranspose2d(in_channels=32, num_filter=32, kernel=3, stride=2, pad=1, output_padding=1)
–
GeLU()
–
ConvTranspose2d(in_channels=32, num_filter=1, kernel=3, stride=2, pad=1, output_padding=1)
–
sigmoid()
- •
Transition function: Here, we consider 8 different MLPs, each of which has the following architecture:
–
Linear(in_features=17, out_features=64, bias=True)
–
LeakyReLU(negative_slope=0.2)
–
Linear(in_features=64, out_features=64, bias=True)
–
LeakyReLU(negative_slope=0.2)
–
Linear(in_features=64, out_features=1, bias=True))
- •
Normalizing Flow: As the nonstationary prior for the noise variables, we use $1D$ conditional normalizing flows, which are 1-layer neural spline flows conditioned on the auxiliary variable $\mathbf{u}$. Before taken as, the auxiliary variable $\mathbf{u}$ is embedded. This is done by a single linear layer with 32 dimensions in the synthetic experiments, and an MLP with the following architecture in the cartpole experiment:
–
Linear(in_features=7, out_features=64, bias=True)
–
LeakyReLU(negative_slope=0.2)
–
Linear(in_features=64, out_features=32, bias=True)

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Does the identifiability still exist when I have any external input, like the control input in the state-space model?
2. What's the exact definition of $s_t$, and what's the motivation for incorporating $s_t$ into the model?
3. (minor) For Table 1, do you use any stats test to confirm your approach is the best?
soundness: 3
strengths: In Table 3, the author(s) claim, "Our approach is the first to predict future states of a dynamical system by utilizing a transition function that identifies the underlying, true dynamics." As far as I know, many approaches can do similar things without a theoretical foundation for the identification.
summary: This paper proposes a structure with a theoretical foundation for identifying latent dynamics under certain assumptions. The implementation is standard. The experiments' results show various advantages in different scenarios, even when the assumptions are violated.
weaknesses: Perhaps a theoretical or practical time complexity analysis would be more appropriate.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. Why was such a small prediction horizon used for testing? I found this odd for the synthetic data.
2. I see that training sequences used to train the model are relatively short. Is there a reason why?
soundness: 3
strengths: My favorite thing about this paper is the simplicity of the proposed approach. By just simply changing the form of the latent dynamics model, i.e., as opposed to additive noise, which is the standard, the noise is an input to the transition function, and assuming a dimension-dependent dynamics function, one can recover the states and dynamics (under some assumptions). I think this is a massive break through for people who learn state-space models.
summary: In this paper, the authors introduce a theory under which one can identify both the latent states as well as the latent dynamics from observational data (up to permutations and invertible transformations). Next, the authors introduce a practical algorithm, based on sequential variational autoencoders, for learning identifiable state-space models. Experiments provide empirical evidence that proposed approach can correctly recover the latent states and the latent dynamics.
weaknesses: While the paper is strong, there a couple of minor weaknesses. Firstly, is that assumption 2 doesn't seem to align with one of the motivating examples in the introduction, which is model-based RL. Specifically, assumption 2 does not allow for process noise/control values that are temporally correlated. But in model-based RL, the process noise *is* temporally correlated because the control value chosen at time $t$ depends on the previous state of the system. I found this contrast jarring.

Second, the experiments section is great but I think it would be great to have a seqVAE baseline that is more similar to proposed generative model. In the experiments section, the baseline seqVAE is the KalmanVAE, which is a hierarchical SSM. While the proposed approach outperforms KalmanVAE, its not clear whether it outperforms it is truly better or if there is a model mismatch between the true generative model and the one used by KalmanVAE. I think training a seqVAE using the generative model as described in section 2.1 would strengthen the results. 

Also, section 4.4 doesn't make sense at a first glance. If I had the true latents then training a model in a supervised fashion should be better than learning both in an unsupervised fashion. When doing the comparison, were the number of parameters matched? I ask because if your latent state has $K$ dimensions then one would need $K$ neural networks, so it might be the case that the proposed approach is working just because the model has more expressive power.

Lastly, I think the proofs in the appendix could benefit with a little more hand-holding to follow along!

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Is there a reason why the noise posterior $q(\mathbf{s}\_{t} | \mathbf{\tilde x}, \mathbf{z}\_{t-1} )$ explicitly depends on $\mathbf{z}\_{t-1}$? Does this follow from your theory? If so, how?

2. In line 140 the authors write “setting $\mathbf{u}$ to an *observed regime index*...” What do you mean by observed regime index? Do you mean the time index $t$ of the observation $\mathbf{x}(t)$? Or do you mean the observation value $\mathbf{x}(t)$ itself?

3. How is $\mathbf{u}$ used in your practical implementation? Put another way, do you set $\mathbf{u}$ to be an *observed regime index* or $\mathbf{s}\_{t-1}$? Or is it only implicitly define via the prior $p(\mathbf{s}\_t| \mathbf{u})$? In Figure 6, $\mathbf{u}$ is shown as an input to the prior noise model, which is used for forecasting. Can you please explain how you choose $\mathbf{u}$ in each experiment?

4. What does $\mathbf{s}$ represent in the synthetic experiments? Similarly, what does it represent in the Cartpole dataset?

5. In Table 1, why is $\text{MCC}[\mathbf{z}] > \text{MCC}[\mathbf{s}]$ for all models? Is the noise variable harder to infer in practice? 

6. How do you interpret the inferred $\mathbf{s}$ in the MOCAP datasets? 

7. In the synthetic dataset description of the appendix (lines 1235, 1236) the authors write: the number of environments is R = 20. What do these environments represent?

8. Are the sufficient variability assumptions (A3, A4) satisfied in the synthetic and Cartpole experiments? If so, can you please explain why?

9. Why are the simulated trajectories (of the synthetic dataset) so short? Meaning: what is the motivation behind such a setting?
soundness: 3
strengths: - This work tackles a key problem of (nonlinear) dynamical system modelling, namely the identifiability of the transition function underlying the dynamics;
- The authors naturally extend the identifiability theory from previous works to the identification of the nonlinear transition function of the hidden process. Their proofs are well written and easy to follow;
- The authors demonstrate that their model implementation of their theory outperforms different baselines, in different settings. The results are very convincing (see however the questions below).
summary: This work deals with the general problem of identifying (i.e. inferring) the nonlinear transition function underlying the dynamics of complex dynamical systems. The authors extend recent theoretical work on identifiability of latent states in dynamical systems to the identifiability of their transition function. Furthermore, the authors provide a neural-based inference model that implements the key aspects of the theory and relies on neural variational inference. The authors then test their methodology on three settings: (i) a synthetic dataset that mirrors their generative model; (ii) a cartpole simulation which is represented as a video signal; and (iii) three sets of experimental motion capture data. The results demonstrate the proposed models outperform the baselines, not only on the identification of the latent variables, but also on forecasting tasks.
weaknesses: In my view, the main weaknesses of the paper have to do with its presentation. The three most pressing points are:

1. There is somewhat of a disconnect between the theory section and the practical implementations, which I think could easily be improved. Especially when it comes to the experimental section. Questions 2 – 8 below point to concepts or ideas which are not clearly explained, that could be explained better or that are simply absent. In particular, it is not clear to me how the auxiliary ($\mathbf{u}$) and noise ($\mathbf{s}$) variables are used (or represented) in the experiments. To improve the readability of the paper, I'd suggest the authors explicitly describe the role of $\mathbf{u}$ and $\mathbf{s}$ in each experiment.

2. The target datasets could be explained and motivated better. It is not clear from a first reading why these datasets are relevant in the context of the model (see e.g. questions 3 and 4 below). This point is clearly related to point 1 above. I believe that if the authors explain what $\mathbf{u}$ and $\mathbf{s}$ represent in each dataset, it will become clearer why those target datasets are interesting.
 
3. The paper builds on a large collection of previous works which deal with identifiability and go all the way back to independent component analysis (ICA). The authors assumed the reader is versed in the development of ICA, from its linear to its nonlinear versions, and simply present e.g. the augmented dynamics (line 126) and the auxiliary variable (line 130), or the identifiability assumption, without any further explanations and context. There is however some nice physical interpretation underlying e.g. the conditional independence assumption of the noise process given the auxiliary variable, which can be traced back to Packard et al. [1980], and that is also nicely explained (albeit in a different sense) by e.g. Hyvarinen et al. [2019]. Even in the introduction, lines 41-49, the authors could already mention that their work builds on classical ideas of ICA. In short, I think the authors could better exploit the history of nonlinear ICA and its application to dynamical systems to better motivate their proposal. 

Note:  *I will happily increase my score once the authors answer the questions below* and explain how they will try to incorporate this information back into the paper.

**Other comments and/or typos** 
- in line 88 the authors write *the generative function*. Do the authors mean the function $\mathbf{g}$? I’d suggest, for the sake of clarity, to either stick to the “instantaneous mapping” label used later in line 112, or better use the more standard “emission model” label.  
- in your KL loss term (lines 290 inside Algorithm 1) you write that the posterior over the noise $\mathbf{s}$ is conditioned on $\mathbf{u}$. This is a typo, isn’t it?

**References**
- Packard et al. [1980]: Geometry from a time series
- Hyvarinen et al. [2019]: Nonlinear ICA Using Auxiliary Variables and Generalized Contrastive Learning

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Does the proposed framework requires the dimensionality of latent space and noise to be known?
If not, what would be the procedure of estimation?

Does the proposed framework allow for sparse noise?
soundness: 3
strengths: The identifiability was theoretically defined and justified with given assumptions.

The proposed framework was empirically evaluated and compared with existing methods across synthetic and real-world data.
summary: This paper presents a framework for identifying the unknown transition function in latent dynamical systems. The authors emphasize that previous research focused on identifying latent states, and lack the ability to accurately predict future states, particularly in complex nonlinear systems.

The authors establish the theoretical identifiability of both the latent states and the transition function. They leverage the assumption of "sufficient variability," meaning that the latent noise and states exhibit enough variation across different environments to disentangle the underlying dynamics. The framework also relies on the assumption that the augmented transition function is bijective. 

The framework was tested across a range of scenarios, including synthetic datasets, a simulated cartpole environment, and real-world Mocap data. The results demonstrate that the model effectively recovers the true underlying dynamics, and is accurate in predicting future states.

After the discussion, the authors have addressed my concerns. I have decided to raise the score.
weaknesses: The claimed contribution "the first framework that allows for the identification of the unknown transition function alongside latent states and the generative function" seems overstated. There are existing methods to identify latent states, transition and generative functions together, e.g. [Sussillo16, Schimel21, Hess23, Zhao23]. 

The proposed framework is not demonstrated as general as claimed like "... improves generalization and interpretability of dynamical systems...". The authors focused on typical RL and robotics scenario. There exist dynamical systems in other fields do not satisfy those assumptions or are in completely different regime. This could disadvantage the compared methods with weaker assumptions (e.g. NODE).
Though the authors used Mocap dataset to showcase then it violates the identifiability assumptions, the system itself is still of RL and robotics scenario. It is totally valid to narrow down problems to solve them better. However, the statements and claims should be more specific.

Meta-review: This paper provides a theoretical state-space modeling framework for identifying both latent states and nonlinear latent dynamics, in contrast to previous works that assumes linear transitions. It also introduces a practical algorithm based on variational autoencoder and shows superior performance in recovering the latent dynamics, making accurate future prediction and adapting to new environments.

All reviewers appreciate its theoretical foundation to the proposed framework and a practical algorithm. Experiments also show robust performance across diverse datasets and strong comparison to baselines.

Decision: Accept (Poster)