Technical Summary: MetaKoopman

Problem Statement

Modern autonomous systems operate in complex, non-stationary environments where underlying dynamics are neither fixed nor fully known. Factors such as terrain changes (e.g., asphalt to ice), variable payloads, and mechanical degradation induce distributional shifts, creating a mismatch between training conditions and deployment realities. Standard approaches face two primary limitations:

1. Global Model Assumption: Most models assume a single, globally valid dynamics model, which is infeasible for time-varying real-world systems.
2. Adaptation vs. Efficiency Trade-off: While meta-learning offers a path to rapid adaptation, existing frameworks often rely on gradient-based inner loops or iterative optimization. These introduce computational latency and overhead, making them unsuitable for real-time deployment. Furthermore, many existing Koopman-based models are trained offline and remain static at test time, limiting their responsiveness to changing dynamics.

Additionally, safety-critical applications require not only accurate predictions but also principled uncertainty quantification (both epistemic and aleatoric) to assess reliability.

Methodology

The paper proposes MetaKoopman, a Bayesian meta-learning framework that integrates the structured efficiency of Koopman operators with the adaptivity of meta-learning. The core innovation is the use of a Matrix Normal-Inverse Wishart (MNIW) prior over the Koopman operator, enabling closed-form Bayesian updates.

1. Latent Koopman Dynamics Model

The system models nonlinear dynamics $x_{t+1} = f(x_t, u_t)$ by lifting state-action pairs into a finite-dimensional latent space via an embedding function $g$.

• Sequence-Based Embedding: Unlike methods embedding single state-action pairs, MetaKoopman uses a context encoder $G_\theta$ to map a window of past states and actions $[x_{t-q}, \dots, x_t]$ into a latent representation $\tilde{z}_t$. This captures temporal dependencies, supported by Takens' theorem.
• Linear Evolution: In the latent space, dynamics evolve linearly: $\tilde{x}_{t+1} = K\tilde{z}_t$, where $K$ is the Koopman operator.
• Future Control Encoding: Future control inputs are encoded via a dedicated action encoder to ensure consistency with latent dynamics.

2. Closed-Form Bayesian Inference

Instead of estimating a fixed operator, MetaKoopman places a distribution over plausible operators.

• Probabilistic Formulation: The latent dynamics are modeled as $\tilde{x}_{t+1} = K\tilde{z}_t + \epsilon$, where $\epsilon \sim \mathcal{N}(0, \Sigma)$.
• Conjugate Prior: The likelihood of latent transitions is paired with an MNIW prior: $K, \Sigma \sim \text{MNIW}(\breve{M}, \breve{V}, \breve{\nu}, \breve{\Psi})$.
• Analytic Posterior: Given recent trajectory data, the posterior parameters are updated in closed-form (Lemma 4.1), avoiding iterative optimization.
• Predictive Distribution: The framework derives a closed-form posterior predictive distribution for future states (Lemma 4.2), approximated by a multivariate Gaussian, capturing both model and process uncertainty.

3. Meta-Learning the Prior

The parameters of the MNIW prior ($\theta$) are meta-learned across a distribution of tasks (different operating conditions).

• Objective: Minimize the negative log-likelihood of ground truth on held-out query sets after updating the prior with adaptation data (Eq. 15).
• Tempering Mechanism: A scalar $\beta \in (0, 1]$ is meta-learned to rescale the prior precision ($\breve{V}' = \beta^{-1}\breve{V}$). This allows the model to dynamically balance prior knowledge against new evidence, increasing responsiveness to online adaptations while maintaining calibrated uncertainty.

4. Uncertainty-Aware Planning

To integrate MetaKoopman into motion planning, the authors introduce a Variational Action Encoder.

• Cost-Guided Latent Action Sampling (CLAS): Instead of deterministically encoding future actions, the encoder models a distribution over latent actions regularized to match a standard normal prior.
• Efficiency: This allows the planner to sample future action embeddings directly from $\mathcal{N}(0, I)$ without evaluating the encoder for every sample. This accelerates trajectory rollouts by an order of magnitude, enabling efficient sampling-based planning using the linear Koopman dynamics.

Key Contributions

1. Bayesian Meta-Learning of Koopman Operators: The first framework to combine Koopman structure with Bayesian meta-learning, providing an analytic test-time posterior over the operator and a closed-form posterior-predictive distribution for principled uncertainty quantification.
2. Meta-Learned Tempering: A scalar $\beta$ meta-learned during training rescales the prior precision at test time, optimizing the trade-off between prior knowledge and new data for adaptive responsiveness.
3. Variational Action Encoder: A mechanism that accelerates trajectory generation by an order of magnitude, removing the computational bottleneck in model-based planning.
4. Comprehensive Evaluation: Evaluation on a full-scale autonomous truck dataset (snow, ice, mixed friction) and five simulated environments (HalfCheetah, Hopper, Walker, etc.) demonstrating state-of-the-art accuracy and robustness.
5. Real-World Deployment: Successful live trials on a 37.5-ton autonomous truck, demonstrating adaptive braking on ice and safe evasive maneuvers where static models failed.

Experimental Results

The paper evaluates MetaKoopman against strong baselines, including Deep Koopman Operators (DKO), Gradient-Based Adaptive Learners (GrBAL), Bayesian MAML, Bayesian Neural Networks (BNNs), Ensemble models (EMLP), and Neural ODEs.

• Prediction Accuracy: MetaKoopman consistently achieves lower Mean Squared Error (MSE) across all environments, including the real-world truck dataset and simulated tasks with distribution shifts (e.g., varying slopes, gravity, friction, and joint failures).
• Adaptation Efficacy: Analysis of the truck dataset shows that after Bayesian meta-updates using recent trajectory data, the distribution of prediction errors shifts significantly toward lower values, indicating successful alignment with evolving dynamics.
• Real-World Performance:
  • Precise Stops on Ice: In scenarios involving a transition from snow to polished ice, non-adaptive models underestimated stopping distances and overshot targets. MetaKoopman adapted online to the reduced friction, enabling accurate braking.
  • Evasive Maneuvers: When a virtual obstacle appeared on a low-traction surface, non-adaptive models attempted infeasible braking. MetaKoopman detected the distribution shift via increased prediction error, adapted its model, and successfully executed a safe lane change.

Significance and Claims

The paper claims that MetaKoopman bridges the gap between structured modeling and meta-learning, offering a tractable, uncertainty-aware framework for online adaptation under distributional shifts.

• Practical Viability: The authors demonstrate that closed-form Bayesian meta-learning is not only theoretically sound but viable for safety-critical, real-time systems (e.g., heavy-duty autonomous trucks).
• Design Philosophy: The framework deliberately models uncertainty over the Koopman operator while keeping the encoder deterministic. The authors argue that in structured environments, the dominant source of variability arises from changes in system dynamics rather than the representation itself. This focus allows for data-efficient updates where they matter most.
• Future Impact: The results suggest a practical path toward principled uncertainty in control, potentially extending to reinforcement learning, multi-agent systems, and partially observable environments.

The work is supported by Vinnova, Sweden, through the AllDrive project, with field experiments conducted in northern winter conditions.