Switching Koopman Latent World Model

Integrate this at the world-model component of a discrete-action model-based RL agent, or immediately after a visual encoder in a video-prediction model. The encoder maps an observation batch to $z_t=\phi_\theta(o_t)\in\mathbb{R}^d$; maintain one trainable matrix $K_a$ for every action $a\in\{1,\ldots,A\}$. For a replay sequence $(o_t,u_t,o_{t+1},\ldots,o_{t+H})$, compute $z_t=\phi_\theta(o_t)$ and recursively set z_pred = K[u[t]] @ z_pred for each step. Minimize the sum of squared errors between each predicted latent and the encoded target latent, optionally add a decoder reconstruction loss, and add the spectral-norm penalty from the formula above. Pseudocode: z = encoder(o[t]); loss = 0; for h in 1..H: z = K[u[t+h-1]] @ z; loss += mse(z, encoder(o[t+h])); for a in actions: loss += lambda_K * relu(power_spectral_norm(K[a])-tau)**2; backprop(loss). The paper supplies the switching linear-operator structure; the neural encoder, finite rollout loss, and spectral regularizer are the adaptation. Estimate each $\|K_a\|_2$ with one or two power iterations per update rather than an expensive SVD; optionally estimate $\rho(K_a)$ offline from eigenvalues. Start with CartPole or Acrobot using an MLP encoder and compare against a standard nonlinear MLP one-step world model at equal latent dimension and parameter count. Evaluate one-step and 20-step prediction error, model-based planning reward, gradient stability, and rollout divergence. A successful signal is substantially lower long-horizon latent error or higher planning reward at equal rollout FLOPs, with fewer exploding imagined trajectories.