Identifying latent state transition in non-linear dynamical systems

TL;DR

The paper tackles identifying latent states and the nonlinear transition function of dynamical systems from high-dimensional observations by introducing an identifiable state-space framework with augmented dynamics and an auxiliary variable to induce nonstationarity or autocorrelation. It proves identifiability for both latent states (Theorem 1) and the transition function (Theorem 2) under structured assumptions, and implements a variational autoencoder-based algorithm that jointly recovers latents, dynamics, and process noise. Through synthetic and modified Cartpole experiments, the method achieves superior latent recovery, calibrated uncertainty, and longer-horizon prediction, while demonstrating data-efficient adaptation to new environments. The work advances interpretable, transferable dynamical modeling with potential applications in model-based control and policy learning, while acknowledging limitations tied to the identifiability assumptions and simulated testbeds.

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.

Figures (7)

• 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 1yao2021learning. In addition to this, our main contribution is the identifiability of the transition function ensured by Theorem 2.
• 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 \ref{['fig:app:uncertainty']} for all regimes.
• Figure 3: Mcc vs. Mse results for ablations for different time steps.
• Figure 4: A comparison of MSEs achieved by our approach and two-stage LEAP training.
• Figure 5: Extended version of Figure \ref{['fig:uncertainty']}