Data-Efficient Koopman Operator Discovery for Soft Robots using Kolmogorov-Arnold Networks (KANs)

This repository contains the implementation and experimental code for the paper "Data-Efficient Koopman Operator Discovery for Soft Robots using Kolmogorov-Arnold Networks".

📌 Project Overview

Soft robots offer flexibility and safety advantages over rigid counterparts but are notoriously difficult to control due to their infinite degrees of freedom and high nonlinearity. Traditional data-driven methods (like MLPs) often function as "black boxes" and require prohibitive amounts of data.

This project proposes a novel modeling framework that integrates Kolmogorov-Arnold Networks (KANs) into Koopman Operator Theory. By replacing fixed polynomial basis functions with learnable splines on network edges, we aim to discover a lifting function $\psi(x)$ that linearizes the dynamics while remaining data-efficient and interpretable.

Key Hypotheses Verified

1. Data Efficiency: KANs converge to asymptotic accuracy using significantly fewer training trajectories than polynomial baselines (EDMD).
2. Compactness: KANs achieve accurate linearization with a lower latent dimension (fewer observables) compared to fixed dictionaries.
3. Real-Time Viability: The inference time remains low (~2ms), making it suitable for real-time control loops (>50Hz).

---

📂 Repository Structure

🚀 Quick Start

• mini_project.ipynb: The main entry point. An example notebook that trains both the KAN-Koopman and Polynomial-Koopman models, compares their performance, and visualizes the results. Start here.
• kan_koopman.py: Implementation of the KAN-based Koopman operator.
• poly_koopman.py: Implementation of the baseline Polynomial EDMD model.

🧪 Experiments (scripts/)

The scripts/ directory contains the drivers for the experiments detailed in the paper:

• limited_trajectories.py: Tests Hypothesis 1 (Data Efficiency). Sweeps over the number of training trajectories ($N_{traj} \in [2, 13]$) to measure the reduction in data required.
• variable_observables.py: Tests Hypothesis 2 (Compactness). Sweeps over the number of latent observables to determine the sensitivity of the model to dimension size.
• curric_learning_effect.py: Evaluates the impact of curriculum learning on training stability and preventing convergence to suboptimal local minima.
• benchmark_inference.py: Measures the wall-clock inference time on CPU to verify real-time control feasibility.
• utils.py: Helper functions for data loading and processing.

📊 Data & Results

• data/: Contains the training dataset from the pneumatic soft robot arm, adapted from Bruder et al. (2019). The data consists of state-action snapshot pairs recorded at 10Hz.
• experiments/: Raw logs and data generated by the scripts.
• figures/: Output directory for plots.
• create_plots.ipynb: Notebook used to generate the specific figures for the final report.

---

⚙️ Methodology

The Koopman Framework

We seek a lifting function $\psi: \mathcal{M} \to \mathbb{R}^N$ such that the nonlinear dynamics $x_{k+1} = f(x_k, u_k)$ become linear in the lifted space: $$ \psi(x_{k+1}) = A \psi(x_k) + B u_k $$

Models Implemented

1. Polynomial Baseline (poly_koopman.py): Uses a fixed dictionary of polynomial basis functions (standard EDMD). Vulnerable to the "curse of dimensionality".
2. KAN-Koopman (kan_koopman.py): Uses a Kolmogorov-Arnold Network to learn the observables $\psi(x)$. KANs place learnable univariate B-spline activation functions on edges, offering high expressivity with fewer parameters.

---

📈 Results Summary

• Data Efficiency: The KAN model approaches its performance floor with $<50%$ of the data required by the polynomial baseline.
• Low Dimensionality: Peak performance was achieved with only 2 learned observables, whereas polynomials required higher orders to compete.
• Speed: Inference time is approximately 2ms on a single CPU thread, validating utility for inertial control.

---

📚 References

This work utilizes data and concepts from:

• D. Bruder, B. Gillespie, C. D. Remy, and R. Vasudevan, "Modeling and Control of Soft Robots Using the Koopman Operator and Model Predictive Control," RSS, 2019.
• Z. Liu et al., "Kolmogorov-Arnold Networks," 2024.