Koopman Theory.
Linear thinking
for nonlinear worlds.

Suppose you have a dynamical system evolving in discrete time: $x_{k+1} = f(x_k)$, where $f$ is some messy nonlinear map. The state $x \in \mathbb{R}^n$ follows a curved, unpredictable path. Classical control chases $x$ directly, and pays a steep price when $f$ becomes stiff, multi-modal, or ill-conditioned.

Koopman inverted the question. Rather than tracking the state, track functions of the state — called observables. An observable is any scalar quantity you can measure from $x$: its energy $\|x\|^2$, its components $x_1, x_2$, some nonlinear combination $\sin(x_1) x_2$, anything. The key fact: the evolution of observables is linear, even when the state is not.

The catch — and it is a real one — is that this linear operator $\mathcal{K}$ acts on an infinite-dimensional space of functions. To make the idea useful we need finite approximations: pick a good enough dictionary of observables, project the operator down, and pray the projection captures the dynamics you care about. Much of modern Koopman work is about that single engineering choice.

Linearity is not a property of dynamics, but of the space in which you choose to describe them. — the Koopman perspective, in one line

Let $f: \mathcal{M} \to \mathcal{M}$ be a (possibly nonlinear) discrete-time dynamics map on some state space $\mathcal{M}$. Let $g: \mathcal{M} \to \mathbb{C}$ be an observable — a scalar-valued function of the state. The Koopman operator $\mathcal{K}$ is defined by how it acts on $g$:

That's it. The entire theory hangs off that one line. Two consequences deserve immediate attention:

1. Linearity is automatic. For any observables $g, h$ and scalars $\alpha, \beta$: $\mathcal{K}(\alpha g + \beta h) = \alpha g \circ f + \beta h \circ f = \alpha (\mathcal{K} g) + \beta (\mathcal{K} h)$. The underlying $f$ can be as vicious as you like; $\mathcal{K}$ doesn't care.
2. Infinite-dimensional, usually. The space of all observables on $\mathcal{M}$ is an infinite-dimensional function space (typically $L^2(\mathcal{M})$). So we've swapped a finite-dim nonlinear problem for an infinite-dim linear one — a classic functional-analytic trade.

Eigenfunctions: the secret ingredient

A Koopman eigenfunction is an observable $\varphi$ with a very special property: $\mathcal{K}\varphi = \lambda \varphi$. Along any trajectory, $\varphi$ evolves by pure scalar multiplication — its geometry is trivial. If you can find enough eigenfunctions to span your state of interest, you have a full change of coordinates that diagonalises the dynamics. This is the nonlinear analogue of finding the modes of a linear system.

Consider the two-dimensional nonlinear system (a classic pedagogical example, due to Brunton):

Define the lifted state $z = \big[x_1,\ x_2,\ x_1^2\big]^\top \in \mathbb{R}^3$. Differentiate and substitute:

Below, watch a trajectory run in the original 2D state space (left) and in the lifted 3D observable space (right). The left portrait is curved; the right one is a clean linear flow on a plane.

The curved line on the left and the straight-in-a-tilted-plane line on the right describe the same physical trajectory. The only difference is the coordinate chart. This is the Koopman promise in miniature. The bad news: finding the right observables analytically is rarely possible. The good news: you can often learn them from data.

Dynamic Mode Decomposition (DMD), introduced by Schmid in 2010 and later tied to Koopman theory by Rowley & Mezić, is the workhorse. Strip away the mystique and DMD is just this: given snapshots of a system, find the best linear operator that maps each snapshot to the next one.

The basic DMD recipe

1. Collect snapshots. Arrange $m$ consecutive state measurements as columns: $X = [x_0, x_1, \dots, x_{m-1}]$ and $X' = [x_1, x_2, \dots, x_m]$.
2. Solve the least-squares problem. Find $A$ minimising $\|X' - A X\|_F$. The closed-form answer is $A = X'\, X^{+}$, where $X^+$ is the pseudoinverse.
3. Extract modes. Eigendecompose $A = W \Lambda W^{-1}$. Each eigenvector of $A$ is a DMD mode; the eigenvalue gives its frequency and growth/decay rate.
4. Predict. Future states are a linear combination of modes, each evolving as $\lambda_i^k$. You have a forecast.

DMD applied to raw state data only captures the part of the Koopman operator that happens to be linear in $x$. For truly nonlinear systems you need Extended DMD (EDMD): lift each snapshot through a dictionary of observables first, then run DMD in that lifted space.

A minimal Python implementation

This is the entire EDMD algorithm. No library required beyond NumPy:

import numpy as np

def dictionary(x):
    # polynomial observables up to order 2, including constant
    return np.vstack([np.ones_like(x[0]),
                      x[0], x[1],
                      x[0]**2, x[0]*x[1], x[1]**2])

def edmd(X, X_next):
    Psi  = dictionary(X)         # N_dict × m
    PsiP = dictionary(X_next)    # N_dict × m
    K    = PsiP @ np.linalg.pinv(Psi)   # N_dict × N_dict
    return K

def predict(K, x0, steps):
    z = dictionary(x0.reshape(2,1))
    traj = [x0.copy()]
    for _ in range(steps):
        z = K @ z
        traj.append(np.array([z[1,0], z[2,0]]))  # recover x from dict
    return np.array(traj)

Below: Brunton's slow-manifold system — a 2D nonlinear system with an exact finite-dimensional Koopman embedding on $\{1, x, y, x^2\}$. With the right four terms, the linear prediction (dashed) and the true nonlinear flow (solid) overlap almost perfectly. Drop $x^2$ from the dictionary and watch prediction quality collapse. The dictionary choice is everything.

Model Predictive Control works beautifully when the plant is linear: the optimisation over a horizon reduces to a convex Quadratic Program (QP), solvable reliably in a few milliseconds. For nonlinear plants you face an NLP — slower, non-convex, and dependent on a good initial guess. Koopman offers a tidy escape:

1. Collect trajectory data from the plant under excitation (system ID).
2. Fit a Koopman model with control (Koopman-with-inputs, or KIC): $z_{k+1} = A z_k + B u_k$, where $z = \Psi(x)$.
3. Solve linear MPC on $(A,B)$ over the lifted state. Standard QP machinery — OSQP, qpOASES, HPIPM — applies directly.
4. Apply the first input, measure, re-lift, repeat.

Where it shines (and where it doesn't)

A domain example: electric motor control

PMSM torque control is an instructive case. The cross-coupled $dq$ voltage equations are mildly nonlinear due to $\omega L i$ terms; field-weakening and saturation make things worse. A polynomial-dictionary Koopman model fit from dyno data can replace a gain-scheduled PI + feed-forward controller with a single linear MPC that handles saturation constraints explicitly — without the real-time cost of a full nonlinear solve. The same principle scales to axial-flux machines where cross-coupling is stronger still.

Reinforcement Learning from raw pixels or high-dimensional sensors is hard partly because the learned world model has to be both expressive and tractable for planning. Koopman-flavoured ideas offer a natural compromise: learn a nonlinear encoder $\Psi_\theta$ and a linear dynamics matrix $K$ in the latent space.

Why this helps an RL agent

• Planning is cheap. Rolling out $n$ steps in latent space is $n$ matrix-vector products. Monte Carlo tree search and cross-entropy methods become dramatically faster than rolling out a deep nonlinear dynamics network.
• Value functions are polynomial. If the latent dynamics are linear and the reward is quadratic in the latent, Bellman updates have closed-form solutions à la Riccati. Critic networks become optional.
• Stability is analysable. Eigenvalues of $K$ tell you whether the learned model will blow up over long horizons. You can regularise toward $\|K\|_2 \le 1$ during training.
• Composition is linear. Two consecutive actions become $K_{u_2} K_{u_1}$, a matrix product. This opens clean options-frameworks and hierarchical RL.

Representative work

Lusch, Kutz & Brunton (2018) introduced Koopman autoencoders for discovering linear embeddings of nonlinear dynamics. Li et al. (2020, "Learning Compositional Koopman Operators") used them for object-centric physics models in RL. More recent work on Koopman for latent-space planning in pixel-based RL pairs the linear dynamics with learned control matrices $B_\theta(u)$, yielding a bilinear model — close enough to linear for efficient planning, flexible enough to capture action-conditioned behaviour.

A damped, driven pendulum — the hello-world of nonlinear dynamics. The plant below is simulated with full nonlinear physics: $\ddot{\theta} + c \dot{\theta} + \sin\theta = u$. You pick the dictionary of observables; EDMD fits a Koopman model from a short burst of training data; we then roll the linear model forward and compare against the nonlinear truth over a fresh trajectory.

Watch what happens when you remove $\sin\theta$ from the dictionary. The model still looks reasonable short-term, then diverges — a textbook reminder that Koopman fidelity lives or dies by the dictionary.