Contraction-Regularized Latent Dynamics

Integrate the method into a recurrent or transformer-based latent world model immediately after its latent transition function. Let an encoder produce \(z_t\), let the model predict \(\hat z_{t+1}=f_\theta(z_t,a_t)\), and add a small metric network \(L_\psi(z)\) whose diagonal is softplus, giving \(M_\psi(z)=L_\psi(z)L_\psi(z)^\top+\epsilon I\) with \(\epsilon=10^{-4}\). For each minibatch sequence, compute \(J_t=\partial f_\theta(z_t,a_t)/\partial z_t\) using automatic differentiation; initially use a 4--16 dimensional latent state so the full Jacobian is affordable. Pseudocode: encode \(z_{0:H}\); predict \(\hat z_{t+1}\); compute \(M_t=M_\psi(z_t)\) and \(M_{t+1}=M_\psi(\hat z_{t+1})\); form \(A_t=J_t^\top M_{t+1}J_t-\rho^2M_t\); eigendecompose the symmetric part of \(A_t\), retain positive eigenvalues, and add \(\lambda_c\|[A_t]_+\|_F^2\) to the prediction loss. Set \(\rho\) between 0.90 and 0.99 and warm up \(\lambda_c\) for 10% of training. Optionally select two latent endpoints, evaluate Chebyshev basis values \(\mathbf T(s_\ell)\) and derivatives \(\mathbf T_d(s_\ell)\), and solve the paper's endpoint-constrained metric-energy problem every 100 updates as a diagnostic or auxiliary loss. The formulas determine the metric construction, contraction residual, and geodesic energy; Jacobians and largest eigenvalues are estimated empirically by autodiff and eigendecomposition. First test on Pendulum or CartPole with a GRU/MLP world model, comparing the baseline against the regularized model under Gaussian observation noise, action perturbations, and 50-step imagined rollouts. Success requires lower long-horizon latent error, fewer exploding trajectories, and improved model-based control reward at equal parameter count; also report wall-clock overhead and reject the method if one-step accuracy improves while long-horizon stability does not.