Wasserstein-Controlled Gaussian-Mixture Rollouts

1. Integration point: attach a probabilistic transition head to a compact neural state-space model. The head predicts a deterministic mean transition f_theta(z,a), its Jacobian J, and a fixed or learned bank of disturbance modes (pi_j,b_j,Q_j). During planning or inference, maintain tuples (w_i,mu_i,Sigma_i) rather than one latent mean and covariance. 2. Pseudocode: initialize components from the encoder posterior; at each horizon step, for every component i and disturbance mode j, compute mu' = f_theta(mu_i,a_t)+b_j, Sigma' = J_i Sigma_i J_i^T+Q_j, and w' = w_i pi_j; normalize weights; then merge nearby components until the estimated Wasserstein error budget reaches epsilon_W. Evaluate expected rewards and affine chance constraints analytically. 3. Computed versus estimated: mixture moments, quadratic costs, and affine constraint probabilities are analytic; J is obtained by autodiff; exact Wasserstein distance for mixtures is expensive, so use a Sinkhorn transport estimate between components and validate it with Monte Carlo samples. 4. First cheap experiment: train a neural model on a two-dimensional nonlinear oscillator with two disturbance modes and compare single-Gaussian, particle, and eight-component-mixture rollouts for 20 steps. Sweep epsilon_W. The predicted signature is preserved bimodality for the mixture, calibrated chance probabilities, and a sharp increase in rollout Wasserstein error when epsilon_W becomes large enough to merge distinct modes. Test whether the empirical error stays within 20 percent of the estimated Wasserstein budget before that transition.