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

1. Overview

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. Dreamer-style actor-critic updates depend on recursively imagined trajectories, so transition errors compound with horizon and bias return targets.

Koopman Dreamer replaces DreamerV3's deterministic latent dynamics core with a spectrally constrained Koopman backbone. Inspired by Koopman operator theory, it uses 2D rotation-scaling blocks (bounded radii) to represent damping, rotation, and near-periodic modes; linear + low-rank bilinear action terms capture global and state-dependent control effects; stochastic-state modulation supplies local correction. Posterior-conditioned EMA teacher targets + one-step consistency + multi-step rollout + open-loop observation prediction objectives reduce train-imagination mismatch.

2. Problem Formulation

Partially observable continuous control. Complete latent state $y_{t}=(\phi_{t},s_{t})$, where $\phi_{t}$ is the deterministic Koopman evolution state and $s_{t}$ is the discrete stochastic state ($N$ groups of $K$-way one-hot). The deterministic component carries recursively propagatable historical memory; the stochastic component absorbs current-observation correction and unmodeled local factors.

Transition: Koopman backbone produces $\phi_{t+1}=F_{K}(\phi_{t},s_{t},a_{t})$, prior network predicts $p_{\theta}(s_{t+1}\mid\phi_{t+1})$:

$$p_{\theta}(y_{t+1}\mid y_{t},a_{t})=\delta\left(\phi_{t+1}-F_{K}(\phi_{t},s_{t},a_{t})\right)p_{\theta}(s_{t+1}\mid\phi_{t+1})$$

Posterior $q_{\theta}(s_{t}\mid\phi_{t},e_{t})$ only corrects the stochastic state, not $\phi_{t}$ directly.

3. Method

3.1 Koopman Latent Transition

The transition propagates $\phi_{t}$ through spectrally structured autonomous update + action-dependent and stochastic corrections:

$$F_{K}(\phi_{t},s_{t},a_{t})=\operatorname{clip}_{c_{\phi}}\left(A_{K}\phi_{t}+B_{a}\bar{a}_{t}+H_{\theta}(\phi_{t},\bar{a}_{t})+B_{z}z_{t}\right)$$

Decomposed into four additive components: $u_{t}^{K}=A_{K}\phi_{t}$ (autonomous spectral evolution), $u_{t}^{a}=B_{a}\bar{a}_{t}$ (linear action), $u_{t}^{b}=H_{\theta}(\phi_{t},\bar{a}_{t})$ (state-dependent action modulation), $u_{t}^{z}=B_{z}z_{t}$ (local correction). $\operatorname{clip}_{c_{\phi}}(x)=c_{\phi}\tanh(x/c_{\phi})$ bounds deterministic-state magnitude.

Bilinear state-action interaction (low-rank):

$$H_{\theta}(\phi_{t},\bar{a}_{t})=\beta_{b}W_{o}^{b}\left[\left(W_{\phi}^{b}\phi_{t}\right)\odot\left(W_{a}^{b}\bar{a}_{t}\right)\right]$$

3.2 Spectrally Constrained Koopman Backbone

$A_{K}$ parameterized by 2D rotation-scaling blocks. Deterministic dim $D=2M$, $i$-th subspace:

$$\begin{bmatrix}(A_{K}\phi_{t})^{(2i)}\\(A_{K}\phi_{t})^{(2i+1)}\end{bmatrix}=\rho_{i}\begin{bmatrix}\cos\theta_{i}&-\sin\theta_{i}\\\sin\theta_{i}&\cos\theta_{i}\end{bmatrix}\begin{bmatrix}\phi_{t}^{(2i)}\\\phi_{t}^{(2i+1)}\end{bmatrix}$$

Modal radius $\rho_{i}=\rho_{\min}+(\rho_{\max}-\rho_{\min})\sigma(\alpha_{i})\in[\rho_{\min},\rho_{\max}]$, phase $\theta_{i}=\pi\tanh(\omega_{K,i})$. $A_{K}$ is block-diagonal normal operator: $\|A_{K}\|_{2}=\rho(A_{K})=\max_i\rho_i\leq\rho_{\max}$.

3.3 Posterior-Conditioned Teacher

Corrected teacher trajectory from posterior stochastic state:

$$\bar{\phi}_{t}=\operatorname{clip}_{c_{\phi}}\left(\phi_{t}+\beta_{\phi}\tanh\left(W_{\phi}\operatorname{flat}(s_{t})\right)\right),\quad \bar{z}_{t}=W_{z}\operatorname{flat}(s_{t})$$

EMA target projections reduce drift. Training objectives: one-step Koopman consistency, multi-step teacher/prior rollouts, open-loop observation prediction.

4. Architecture

flowchart TB
    OBS[Observation o_t] --> ENC[Encoder E_θ]
    ENC --> POST[Posterior q_θ
Correct stochastic s_t]
    PREV[Prev state φ_{t-1}, s_{t-1}, a_{t-1}] --> KK[Koopman Backbone A_K
2D Rotation-Scaling Blocks
Spectral radius ≤ ρ_max]
    KK --> PHI[Deterministic state φ_t]
    POST --> PHI
    PHI --> PRIOR[Prior p_θ
Predict s_{t+1}]
    PRIOR --> NEXT[Next state y_{t+1}]
    NEXT --> HEADS[Prediction Heads
Obs/Reward/Continue]
    NEXT --> AC[Actor-Critic
Policy Optimization]
    NEXT --> TEACHER[EMA Teacher
Posterior-Conditioned Targets]
    TEACHER --> KK
    style KK fill:#4f9eff,color:#fff
    style POST fill:#ff6b6b,color:#fff
    style NEXT fill:#51cf66,color:#fff

5. Theoretical Analysis

Derives a multi-step rollout-error bound for the controlled Koopman transition, separating autonomous operator amplification from controlled-interaction effects and additive stochastic/modeling residuals. The spectral radius controls the trade-off between error accumulation and persistent dynamics retention — excessive contraction loses control-relevant persistent information.

6. Experimental Results

6.1 DMC Proprioceptive Continuous Control

9 proprioceptive tasks, 500K steps. Koopman Dreamer achieves best final score on 6 tasks, exceeds DreamerV3 on 8. Largest gains: Acrobot (292.3 vs 131.7) and Hopper Stand (859.1 vs 650.4).

6.2 Open-Loop Prediction

32-step context, 64-step rollout without future observations. Deterministic latent MSE reduced 89.4% (9/9 tasks), proprioceptive MSE reduced 23.2% (8/9 tasks), reward MSE 0.929→0.732 (8/9).

6.3 UAV-LiDAR Autonomous Navigation

Forest: 50m map, 120 obstacles, 280-dim 7-ring LiDAR. 240 closed-loop episodes. Koopman Dreamer success 73.8% (vs DreamerV3 53.8%, D4PG 15.8%), lowest failure 26.2%, median final distance 0.72.

6.4 Spectral Radius Sensitivity

6 variants $\rho_{\max}$ from 0.90 to 1.20. Smallest radius 0.849 gives lowest latent MSE but not lowest observable errors — strong contraction makes latent trajectories numerically similar but loses predictive information. Optimal spectral range requires combining open-loop accuracy with closed-loop performance.

7. Discussion and Conclusion

Koopman Dreamer replaces DreamerV3's RSSM deterministic update with a spectrally constrained Koopman backbone, making contraction, rotation, and near-periodic modes part of transition design. Linear + bilinear action terms and stochastic modulation adapt to controlled nonlinear dynamics. Posterior-conditioned teacher and prior-rollout objectives enable the same structured transition to work under both observation-corrected training and posterior-free imagination. A multi-step error bound theory clarifies how the spectral radius balances error accumulation with persistent information retention.

Practical lesson: the most contractive spectral setting is not optimal — practitioners should select the spectral range using decoded observation-prediction accuracy together with closed-loop performance. Limitation: validation is simulation-based, but the structured transition could be adapted to other control/planning frameworks relying on recursive multi-step prediction.