KALIKO: Kalman-Implicit Koopman Operator Learning For Prediction of Nonlinear Dynamical Systems

TL;DR

The paper tackles long-horizon prediction for nonlinear, high-dimensional dynamical systems by reframing embedding learning as a Kalman-filter-based inference problem. KALIKO implicitly learns Koopman embeddings with globally linear latent dynamics through end-to-end training of a Kalman filter and a decoder, avoiding explicit encoder design. It demonstrates faithful reconstructions, interpretable eigenfunctions, and superior open-loop prediction, along with near-oracle closed-loop control in a challenging wave disturbance scenario. This approach fuses Bayesian state estimation with Koopman theory to produce robust, interpretable representations that enhance predictive control in complex environments.

Abstract

Long-horizon dynamical prediction is fundamental in robotics and control, underpinning canonical methods like model predictive control. Yet, many systems and disturbance phenomena are difficult to model due to effects like nonlinearity, chaos, and high-dimensionality. Koopman theory addresses this by modeling the linear evolution of embeddings of the state under an infinite-dimensional linear operator that can be approximated with a suitable finite basis of embedding functions, effectively trading model nonlinearity for representational complexity. However, explicitly computing a good choice of basis is nontrivial, and poor choices may cause inaccurate forecasts or overfitting. To address this, we present Kalman-Implicit Koopman Operator (KALIKO) Learning, a method that leverages the Kalman filter to implicitly learn embeddings corresponding to latent dynamics without requiring an explicit encoder. KALIKO produces interpretable representations consistent with both theory and prior works, yielding high-quality reconstructions and inducing a globally linear latent dynamics. Evaluated on wave data generated by a high-dimensional PDE, KALIKO surpasses several baselines in open-loop prediction and in a demanding closed-loop simulated control task: stabilizing an underactuated manipulator's payload by predicting and compensating for strong wave disturbances.

Figures (7)

• Figure 1: (A) Most methods for learning the linear Koopman dynamics $\bm{K}_\theta$ are explicit, parameterizing an encoder $\bm{g}_\phi$ that maps data $\bm{x}$ to high-dimensional Koopman embeddings $\bm{z}$. Examples include neural networks, fixed dictionaries of basis functions, or learnable dictionaries. (B) KALIKO (ours) instead implicitly learns Koopman embeddings by letting $\bm{g}$ be a Kalman filter and smoother, which is composed of only two models: the latent dynamics $\bm{K}_\theta$ and decoder $\bm{h}_\psi$. We show that these implicitly-recovered representations are highly interpretable and yield strong open-loop prediction and closed-loop control performance.
• Figure 2: (A) KALIKO Training. (i) Filter for $T$ steps, outputting a sequence of filtered distributions $\{(\bm{\mu}_{t \mid t-1}, \bm{\Sigma}_{t \mid t-1})\}_{t=1}^T$ used for the loss. (ii) Run a backward smoother from the final belief $(\bm{\mu}_{T \mid T}, \bm{\Sigma}_{T \mid T})$, yielding the posterior initial belief $(\bm{\mu}_{0\mid T}, \bm{\Sigma}_{0\mid T})$. (iii) Roll this belief forward for $T$ steps to get predicted beliefs $\{(\bar{\mu}_t,\bar{\Sigma}_t)\}_{t=1}^T$. Filtered and predicted means are decoded to observations and the model is trained end-to-end with a reconstruction loss. Learnable modules/parameters and their operations are highlighted in blue. (B) Belief Propagation. The filter's latent belief becomes more certain with more steps. The smoother reuses future evidence to calibrate the initial belief, seeding a strong prediction rollout during training. (C) KALIKO Inference. KALIKO filters on $T_\text{in}$ data points then rolls out its latent dynamics for $T_\text{out}$ steps. The $T_\text{out}$ predicted latents are decoded into the final predicted trajectory.
• Figure 3: KALIKO Reconstructions. KALIKO closely reconstructs the vector field of nonlinear systems without an encoder using globally linear latent dynamics. Shown are some ground truth trajectories in black (solid) and their reconstructions in orange (dashed). A denser heatmap of the reconstruction error is shown in blue (darker denotes more error).
• Figure 4: Limit Cycle Eigenfunction for the VDP System.(A/B) Without parameterizing an encoder, KALIKO implicitly recovers an eigenfunction associated with the Van der Pol system's limit cycle (overlaid in green) with $|\lambda| \approx 1$. (C) Evaluating $\varphi(x)$ along the limit cycle traces out a nearly perfect circle in the complex plane, showing that KALIKO entirely captures the limit cycle dynamics into this eigenfunction. (D) The Koopman mode associated with this eigenvalue corresponds to a vector field that drives trajectories onto the cycle.
• Figure 5: Eigenfunctions for the Undamped and Damped Duffing Oscillator.(A) KALIKO recovers an eigenfunction $\varphi(\bm{x})$ with $\lambda \approx +1$ capturing the invariance of systems with continuous spectra like the undamped Duffing oscillator (UDO). The value of $\varphi(x)$ is nearly constant along each of three distinct cycles (green) with little variation (shown are mean and stdev). Across cycles, the value clearly changes, corresponding to energy invariance. (B) Conversely, KALIKO also reconstructs the damped Duffing oscillator (DDO) with no adjustments. (C/D) KALIKO recovers an eigenfunction capturing the attractive "spiral" behavior about the wells at $(\pm1,0)$, reproducing the results from explicit methods in prior work otto2019_linearlyrecurrentautoencodernetworkslearning.