KEEC: Koopman Embedded Equivariant Control

TL;DR

Koopman Embedded Equivariant Control (KEEC) tackles controlling unknown nonlinear dynamics by learning an embedding that is both equivariant and isometric, ensuring the latent representation faithfully preserves flows, vector fields, and distances. By leveraging the Koopman operator, KEEC achieves linear latent dynamics and derives an analytical greedy control policy from the learned differential information and a latent value function guided by Hamiltonian-Jacobi theory. The approach identifies latent vector fields via the infinitesimal generator and uses a loss that enforces equivariance and metric preservation, yielding stable and efficient control across diverse domains, including image-based Pendulum, Lorenz-63, and wave equation tasks. Empirical results show KEEC outperforms baselines in episodic rewards and trajectory stability, while offering computational advantages through linear latent dynamics and closed-form policy computation.

Abstract

An efficient way to control systems with unknown nonlinear dynamics is to find an appropriate embedding or representation for simplified approximation (e.g. linearization), which facilitates system identification and control synthesis. Nevertheless, there has been a lack of embedding methods that can guarantee (i) embedding the dynamical system comprehensively, including the vector fields (ODE form) of the dynamics, and (ii) preserving the consistency of control effect between the original and latent space. To address these challenges, we propose Koopman Embedded Equivariant Control (KEEC) to learn an embedding of the states and vector fields such that a Koopman operator is approximated as the latent dynamics. Due to the Koopman operator's linearity, learning the latent vector fields of the dynamics becomes simply solving linear equations. Thus in KEEC, the analytical form of the greedy control policy, which is dependent on the learned differential information of the dynamics and value function, is also simplified. Meanwhile, KEEC preserves the effectiveness of the control policy in the latent space by preserving the metric in two spaces. Our algorithm achieves superior performances in the experiments conducted on various control domains, including the image-based Pendulum, Lorenz-63 and the wave equation. The code is available at https://github.com/yyimingucl/Koopman-Embedded-Equivariant-Control.

Figures (14)

• Figure 1: Figure \ref{['fig:pendulum_traj']} is an overview of KEEC. The left panel (Figure \ref{['fig:uncontrolled_vector_field']}) shows the embedded latent uncontrolled vector field of the pendulum, and the right panel (Figure \ref{['fig:controlled_vector_field']}) shows the corresponding embedded latent controlled vector field. Under the controlled vector field, the pendulum contracts to the target point. More information is available in Section \ref{['Sec: experiments']}.
• Figure 2: Qualitative results. In (a-left), the Pendulum started at $(-3,6)$ with a goal state of $(0,0)$. A zoomed-in view of $[-0.1,0.1]\times[-0.3,0.3]$ showed control stability. In (a-right), the goal was the strange attractor $(-8,-8,27)$ of Lorenz-63 system and we visualized the control trajectories of KEEC and the baselines. In (b), we showed the control trajectories of the KEEC in (b-right) and the best baseline SAC in (b-left). This task aimed to steer the system state to the zero state. Control trajectories of other baselines were shown in Appendix \ref{['subAppendix: Description of the tasks']}.
• Figure 3: Quantitative results on evaluation time and ablation studies on latent dimension $n$ and magnitude of the isometric constraint $\lambda_{\text{met}}$. Left: box-plots show the distributions of evaluation time. The white line in the box indicates the median. Our approach is consistently faster than the MPC-based methods and comparable to the RL methods. Right: different dimensions of the latent space (d) and our model’s episodic reward with different magnitudes of $\lambda_{\text{met}}$ (e).
• Figure 4: Comparison of learned value function $V_g$ and latent space with and without(w/o) the $\mathcal{E}_{\text{met}}$ in pendulum task. The colours on the original coordinates and the space indicated the magnitude of $V_g$. The spaces in (b) and (c) were visualized using Locally Linear Embedding roweis2000nonlinear to project from 8 to 3 dimensions.
• Figure 5: The visualization of the learned value function and the KEEC control trajectories of the swing-up pendulum.

Theorems & Definitions (17)

• Theorem 3.1: Equivariant Vector Field
• Proposition 3.2: Equivariant Flow
• Lemma 3.3: Invariant Value Function
• Theorem 3.4: Greedy Policy on Equivariant Vector Fields
• Definition B.1: Manifold, Riemannian Metric and Riemannian Manifold abraham2012manifolds
• Definition B.2
• Definition B.3: Group dructu2018geometric
• Definition B.4: Vector Field
• Definition B.5: Lie Group Action
• Definition B.6: Flow