Distributed Koopman Learning using Partial Trajectories for Control

This paper proposes a distributed data-driven framework for dynamics learning, termed distributed deep Koopman learning using partial trajectories (DDKL-PT). In this framework, each agent in a multi-agent system is assigned a partial trajectory offline and locally approximates the unknown dynamics using a deep neural network within the Koopman operator framework. By exchanging local estimated dynamics rather than training data, agents achieve consensus on a global dynamics model without sharing their private training trajectories. Simulation studies on a surface vehicle demonstrate that DDKL-PT achieves consensus on the learned dynamics, and each agent attains reasonably small approximation errors on the testing dataset. Furthermore, a model predictive control scheme is developed by integrating the learned Koopman dynamics with known kinematic relations. Results on a reference-tracking task indicate that the distributedly learned dynamics are sufficiently accurate for model-based optimal control.

💡 Research Summary

The paper introduces a novel distributed data‑driven framework for learning the dynamics of nonlinear time‑invariant systems, called Distributed Deep Koopman Learning using Partial Trajectories (DDKL‑PT). In multi‑agent systems (MAS), each agent is assigned a short, offline segment of the overall state‑input trajectory and is required to learn a Koopman‑compatible model locally. The local model consists of a deep neural network lifting function g(·,θ) that maps the original state to a higher‑dimensional space, together with linear matrices A, B, and C that describe the lifted‑space dynamics and the reconstruction back to the original state.

The key challenge addressed is that a single agent’s partial trajectory is insufficient to uniquely identify the global dynamics. To overcome this, the authors formulate a multi‑agent optimization problem that minimizes the sum of local loss functions while enforcing consensus constraints on the matrices A, B, C and the lifting parameters θ across all agents. The loss for each agent i captures two errors: (i) the reconstruction error between the true state and C g(x,θ_i) and (ii) the prediction error of the lifted dynamics g(x_{t+1},θ_i) ≈ A_i g(x_t,θ_i) + B_i u_t.

The algorithm proceeds in two stages. In the first stage, with the lifting parameters θ_i fixed, each agent updates its local linear matrices using a novel distributed update rule (equations 12‑13). This rule combines local gradient information with weighted averages of neighboring agents’ auxiliary variables, allowing agents to use different step sizes while still guaranteeing exponential convergence to the optimal matrices. In the second stage, the linear matrices are held constant and the lifting parameters θ_i are refined via a distributed sub‑gradient method (equation 14), where each agent averages the parameters of its neighbors and then takes a gradient step on its own loss. The learning rates are diminishing, ensuring convergence.

The complete procedure is summarized in Algorithm 1: a first loop of S iterations for matrix consensus, followed by a second loop of (\bar{S}) iterations for parameter consensus. After convergence, all agents possess identical estimates ({A^,B^,C^,\theta^}), which constitute a global Koopman model built solely from locally stored data.

Experimental validation is performed on a five‑agent network controlling a surface vehicle with six‑dimensional state (position, orientation, and velocities) and two control inputs (left/right thrust). The overall trajectory from time 0 to 5000 is generated with random inputs; each agent observes a distinct, non‑overlapping segment (e.g., 0‑600, 600‑1600, etc.). The lifting function is a neural network with two hidden layers of 256 ReLU units, mapping the 3‑dimensional velocity vector to an 8‑dimensional lifted space. The DDKL‑PT method is compared against (i) a centralized Deep Koopman Operator (DKO) that uses the full 0‑4000 data, and (ii) a standard multilayer perceptron (MLP) that directly predicts the next velocity from current velocity and control. All methods use the same optimizer (Adam, learning rate 1e‑4), training termination criterion, and are repeated ten times to mitigate randomness.

Results show rapid consensus: the Frobenius‑norm differences of A_i, B_i, C_i, and the parameter vectors θ_i across agents decay to near zero within a few hundred iterations, confirming the theoretical convergence of the distributed updates. In terms of prediction accuracy on the held‑out test interval (4000‑5000), DDKL‑PT achieves an average one‑step error V = 0.0284 ± 0.0041, slightly higher than the centralized DKO (0.0179 ± 0.0016) and the MLP (0.0205 ± 0.0028). The modest performance gap is offset by the benefits of data locality, privacy preservation, and computational scalability.

To demonstrate practical utility, the learned global Koopman model is embedded in a Model Predictive Control (MPC) scheme for the surface vehicle. The MPC solves a finite‑horizon optimal control problem using the linear lifted dynamics together with known kinematic relations. Simulations show successful goal‑tracking and station‑keeping: the vehicle follows a prescribed trajectory and maintains a desired pose with small control effort, indicating that the distributedly learned model is sufficiently accurate for model‑based optimal control.

In summary, the paper makes three major contributions: (1) a distributed Koopman learning algorithm that works with only partial trajectories, (2) a proof‑of‑concept demonstration that consensus on both linear dynamics and nonlinear lifting functions can be achieved without sharing raw data, and (3) an application of the learned model to MPC, confirming its suitability for control tasks. The work opens avenues for further research on asynchronous communication, time‑varying network topologies, privacy‑preserving parameter exchange (e.g., via encryption or federated learning), and real‑time implementation on physical robotic platforms.