Koopman Dreamer: Spectrally Constrained Latent Dynamics for Stable World-Model Imagination

2026-07-21

AI Summary

This work addresses the limitations of existing latent-variable world models, which suffer from mode collapse and error accumulation during long-horizon imagination, thereby compromising sample efficiency and stability in continuous control. To overcome these issues, the authors propose a Dreamer-style world model grounded in Koopman operator theory, featuring a spectrally constrained deterministic latent dynamics core that enables stable long-range predictions. The model explicitly captures damped, rotational, and near-periodic dynamics through a novel two-dimensional rotation–scaling block. It further integrates low-rank bilinear action coupling, stochastic state modulation, and an EMA teacher target, trained with a combination of one-step consistency, multi-step backtracking, and open-loop prediction losses. A theoretical bound on multi-step backtracking error is introduced to disentangle spectral amplification from state mismatch. Experiments on the DeepMind Control Suite and UAV-LiDAR navigation tasks demonstrate substantially improved latent trajectory stability and closed-loop control performance.

Abstract

Latent world models improve sample efficiency in continuous control by optimizing policies over imagined latent trajectories, but common neural transitions offer limited direct control over modal persistence and error accumulation in long rollouts. We propose Koopman Dreamer, a Dreamer-style world model with a spectrally constrained deterministic latent dynamics core. Its Koopman-inspired backbone uses two-dimensional rotation--scaling blocks with bounded radii to represent damping, rotation, and near-periodic modes. Linear and low-rank bilinear action terms capture global and state-dependent control effects, while stochastic-state modulation supplies local correction information. To reduce the mismatch between posterior-conditioned training and prior-only imagination, the model combines posterior-conditioned EMA teacher targets with one-step consistency, multi-step rollout, and open-loop observation-prediction objectives. We further derive a multi-step rollout-error bound that separates amplification by the spectral backbone and bilinear interaction from the additive effects of stochastic-state mismatch and modeling residuals, clarifying the trade-off between error attenuation and long-term information retention. Experimental results on proprioceptive continuous-control tasks from the DeepMind Control Suite and UAV-LiDAR autonomous navigation demonstrate that Koopman Dreamer improves the stability of long-horizon latent rollouts and achieves stronger closed-loop control performance on tasks that rely on high-quality multi-step imagination.

Problem

• latent world models
• error accumulation
• long-horizon imagination
• spectral dynamics
• continuous control

Innovation

• Koopman operator
• spectrally constrained dynamics
• latent world model
• multi-step imagination
• bilinear action modeling