# KEEC: Koopman Embedded Equivariant Control

arXiv: https://arxiv.org/abs/2312.01544 · HTML: https://mlsanity.com/paper/2312.01544/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 38 of its 80 references that are in the map.
- First posted on arXiv 4 Dec 2023.

## Builds on (most cited first, 30 of 38)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Geometric deep learning: going beyond Euclidean data](https://mlsanity.com/paper/1611.08097/) (2016)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Deep learning for universal linear embeddings of nonlinear dynamics](https://mlsanity.com/paper/1712.09707/) (2017)
- [Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning](https://mlsanity.com/paper/1708.02596/) (2017)
- [Modern Koopman Theory for Dynamical Systems](https://mlsanity.com/paper/2102.12086/) (2021)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Visual Reinforcement Learning with Imagined Goals](https://mlsanity.com/paper/1807.04742/) (2018)
- [SO(2)-Equivariant Reinforcement Learning](https://mlsanity.com/paper/2203.04439/) (2022)
- [Learning Compositional Koopman Operators for Model-Based Control](https://mlsanity.com/paper/1910.08264/) (2019)
- [Variational Inference MPC for Bayesian Model-based Reinforcement Learning](https://mlsanity.com/paper/1907.04202/) (2019)
- [Robust Locally-Linear Controllable Embedding](https://mlsanity.com/paper/1710.05373/) (2017)
- [Koopman Q-learning: Offline Reinforcement Learning via Symmetries of Dynamics](https://mlsanity.com/paper/2111.01365/) (2021)
- [Modeling and Control of Soft Robots Using the Koopman Operator and Model Predictive Control](https://mlsanity.com/paper/1902.02827/) (2019)
- [Linear identification of nonlinear systems: A lifting technique based on the Koopman operator](https://mlsanity.com/paper/1605.04457/) (2016)
- [Average-Reward Reinforcement Learning with Trust Region Methods](https://mlsanity.com/paper/2106.03442/) (2021)
- [Prediction, Consistency, Curvature: Representation Learning for Locally-Linear Control](https://mlsanity.com/paper/1909.01506/) (2019)

## Abstract (arXiv)

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.