Research map / Reinforcement learning / Visual and representation RL
Efficient Dynamics Modeling in Interactive Environments with Koopman Theory
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a diagonal, control-affine Koopman linear latent dynamics model that is parallelized over time via FFT-based convolution and whose eigenvalues control gradient flow through time, for long-horizon dynamics modeling, planning (TD-MPC) and model-free RL (SPR).
Problem
Long-horizon prediction of dynamics in interactive (controlled) environments is inaccurate because errors compound, and nonlinear recurrent dynamics models are slow to train sequentially and suffer from unstable gradients through time.
Approach
Linearize controlled dynamics in a learned latent space with a diagonal Koopman matrix plus decoupled action embedding, unroll it in closed form as a circular convolution with a Vandermonde kernel (parallel over time via FFT), and set eigenvalue real parts to control gradient magnitude through time.
Key results
- Over a 100-step horizon, the Koopman model achieves state prediction MSE of 0.007 on hopper-expert and 0.300 on halfcheetah-expert, competitive with GRU (0.006 and 0.280) and better than MLP, Transformer and DSSM on most datasets.
- On halfcheetah-expert training speed at sequence length 500, Koopman reaches 24.56 iterations/s versus 8.91 for GRU and 11.20 for MLP.
- Koopman dynamics in SAC+SPR on DMC pixel tasks outperforms the MLP-based SPR variant for a horizon of 5 while being more efficient (per the paper's text).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use diagonal Koopman dynamics instead of GRU | Offline dynamics modeling · hopper, walker2d, and halfcheetah expert, medium-expert, and full-replay datasets | state prediction error: hopper-expert-v2: 0.006 $\pm$ 0.001 -> 0.007 $\pm$ 0.000; walker2d-expert-v2: 0.100 $\pm$ 0.005 -> 0.090 $\pm$ 0.008; ha | Mixed |
| use diagonal Koopman dynamics instead of GRU | Offline dynamics modeling · hopper, walker2d, and halfcheetah expert, medium-expert, and full-replay datasets | reward prediction error: hopper-expert-v2: 0.001 $\pm$ 0.000 -> 0.001 $\pm$ 0.000; halfcheetah-expert-v2: 0.045 $\pm$ 0.004 -> 0.040 $\pm$ 0.006; | Mixed |
| use diagonal Koopman dynamics instead of DSSM | Dynamics model training · halfcheetah-expert-v2, sequence length 500 | iterations/second: 17.20 -> 24.56 | Helped |
| use constant μ initialization with increasing ω instead of constant μ with random ω (ablation) | Koopman · halfcheetah-expert-v2 | state prediction error: 0.340 $\pm$ 0.047 -> 0.300 $\pm$ 0.060 | Helped |
| use constant μ initialization with increasing ω instead of constant μ with random ω (ablation) | Koopman · halfcheetah-expert-v2 | reward prediction error: 0.041 $\pm$ 0.005 -> 0.040 $\pm$ 0.006 | Helped |
| use constant μ initialization with increasing ω instead of constant μ with random ω (ablation) | Koopman · walker-2d-expert-v2 | state prediction error: 0.104 $\pm$ 0.008 -> 0.090 $\pm$ 0.008 | Helped |
| use constant μ initialization with increasing ω instead of constant μ with random ω (ablation) | Koopman · walker-2d-expert-v2 | reward prediction error: 0.004 $\pm$ 0.001 -> 0.003 $\pm$ 0.000 | Helped |
| use constant μ initialization with increasing ω instead of constant μ with random ω (ablation) | Koopman · hopper-expert-v2 | state prediction error: 0.010 $\pm$ 0.007 -> 0.007 $\pm$ 0.000 | Helped |
| use constant μ initialization with increasing ω instead of constant μ with random ω (ablation) | Koopman · hopper-expert-v2 | reward prediction error: 0.014 $\pm$ 0.005 -> 0.001 $\pm$ 0.000 | Helped |
| increase latent embedding dimension from 128 to 512 (ablation) | Koopman · halfcheetah-expert-v2 | state prediction error: 0.0097 -> 0.0073 | Helped |
| increase latent embedding dimension from 512 to 1024 (ablation) | Koopman · halfcheetah-expert-v2 | state prediction error: 0.0073 -> 0.0075 | Hurt |
| increase prediction horizon from 5 to 20 (ablation) | Koopman model for model-free RL · DeepMind Control Suite | performance | Hurt |
Surprising findings
- The paper notes TD-MPC with the MLP dynamics model suffers training instability at horizon 20, so the vanilla baseline was trained with horizon 5.
- Adding the Koopman consistency loss in both latent and Q-learning spaces was empirically found to promote collapse and unstable training, so it was dropped in model-free RL.
Limitations
- The model is tailored for deterministic environments and does not handle stochastic dynamics.
- Gains in model-based planning and model-free RL are preliminary; the authors hypothesize distribution shift and objective mismatch hurt RL performance.
- Model-free RL results are limited to five DMC environments with a horizon of 5.
Experimental setup
Datasets: D4RL hopper-v2 D4RL walker2d-v2 D4RL halfcheetah-v2 DeepMind Control Suite
Benchmarks: D4RL hopper-v2 D4RL walker2d-v2 D4RL halfcheetah-v2 DeepMind Control Suite (Quadruped Run, Quadruped Walk, Cheetah Run, Acrobot Swingup, Ball in Cup Catch, Cartpole Swingup
Metrics: mean squared error (state and reward prediction) training iterations/second interquartile mean (IQM) optimality gap episodic return
Baselines: MLP dynamics model Transformer dynamics model GRU dynamics model DSSM TD-MPC SAC SPR
Builds on: Koopman theory Deep Koopman TD-MPC SPR diagonal state space models
Closest papers by meaning
- Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models (NeurIPS 2025 · similarity 0.87)
- Simplifying Latent Dynamics with Softly State-Invariant World Models (NeurIPS 2024 · similarity 0.87)
- Reinforcement Learning Under Latent Dynamics: Toward Statistical and Algorithmic Modularity (NeurIPS 2024 · similarity 0.86)
- The Edge-of-Reach Problem in Offline Model-Based Reinforcement Learning (NeurIPS 2024 · similarity 0.86)
- Overcoming Slow Decision Frequencies in Continuous Control: Model-Based Sequence Reinforcement Learning for Model-Free Control (ICLR 2025 · similarity 0.86)
- Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.