# Wasserstein-Controlled Gaussian-Mixture Rollouts

- ID: 2821
- Canonical URL: https://synthcore.org/idea/2821/wasserstein-controlled-gaussian-mixture-rollouts
- API JSON: https://synthcore.org/api/idea/2821.json
- API Markdown: https://synthcore.org/api/idea/2821.md
- Verification status: failed_benchmark
- Source: [arXiv:2608.29272](https://arxiv.org/abs/2608.29272)
- Category: dynamics
- Solves: stability, accuracy, sample-efficiency
- ML areas: world-model, ssm, rl, training-dynamics
- Math tags: probability, optimal-transport, stochastic-processes, statistics, dynamical-systems
- Ratings: usefulness 8/10; difficulty 6/10; novelty 7/10

## Idea description

Replace single-Gaussian uncertainty propagation in a neural state-space or world model with a finite mixture of Gaussian latent states. Each component is propagated through the learned nonlinear dynamics, and components are merged or pruned only when their Wasserstein discrepancy is below a prescribed tolerance, preserving multimodal futures while keeping computation bounded.

## Mathematical statement

A Gaussian mixture is represented as P = sum_i w_i N(mu_i,Sigma_i), where w_i are nonnegative weights summing to one, mu_i is the mean of component i, and Sigma_i is its positive-definite covariance. Approximation quality is measured by the Wasserstein distance W_rho(P,Q) = (inf_gamma E[||x-y||^rho])^(1/rho), where gamma ranges over couplings with marginals P and Q. For a learned transition z_(t+1) = f_theta(z_t,a_t) + epsilon_t, propagate component i and disturbance mode j using mu_(t+1)^(ij) = f_theta(mu_t^i,a_t) + b_j and Sigma_(t+1)^(ij) = J_i Sigma_t^i J_i^T + Q_j, where J_i is the Jacobian of f_theta at mu_t^i, and (b_j,Q_j,pi_j) describe disturbance mode j. The new weight is w_(t+1)^(ij) = w_t^i pi_j. For an affine constraint c^T z <= r, the mixture probability is the weighted sum of Gaussian cumulative distribution functions. Quadratic costs have an exact mixture expectation.

## Key formulas

- $$P=\sum_{i=1}^{m}w_i\mathcal{N}(\mu_i,\Sigma_i),\qquad W_\rho(P,Q)=\left(\inf_{\gamma\in\Gamma(P,Q)}\mathbb{E}_{(x,y)\sim\gamma}[\|x-y\|^\rho]\right)^{1/\rho}.$$
- $$\mu_{t+1}^{ij}=f_\theta(\mu_t^i,a_t)+b_j,\qquad \Sigma_{t+1}^{ij}=J_i\Sigma_t^iJ_i^\top+Q_j,\qquad w_{t+1}^{ij}=w_t^i\pi_j.$$
- $$\Pr(c^\top z\le r)=\sum_iw_i\Phi\left(\frac{r-c^\top\mu_i}{\sqrt{c^\top\Sigma_i c}}\right).$$
- $$\mathbb{E}[z^\top Rz]=\sum_iw_i\left(\mu_i^\top R\mu_i+\operatorname{tr}(R\Sigma_i)\right),\qquad W_2(P,\widehat{P})\le\varepsilon_W.$$

## Implementation notes

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.

## Verification

- Status: failed_benchmark
- Mechanism evidence: yes
- Mechanism confirmed: no
- Practical verdict: harms
- Verdict: Built a readable nonlinear Gaussian-mixture rollout prototype with disturbance-mode branching, analytic chance constraints, moment calculations, and greedy Gaussian-Wasserstein merging. Mode separation matched the predicted 2d exactly, merging transitioned at the measured pair W2 of 1.8, and empirical error rose sharply after merging. The mixture also retained multimodality that the moment-matched Gaussian lost, with chance-probability gaps up to 0.209 and rollout W2 error reaching 1.052 by horizon 8.

### Mechanism check

- Verdict: Built a readable nonlinear Gaussian-mixture rollout prototype with disturbance-mode branching, analytic chance constraints, moment calculations, and greedy Gaussian-Wasserstein merging. Mode separation matched the predicted 2d exactly, merging transitioned at the measured pair W2 of 1.8, and empirical error rose sharply after merging. The mixture also retained multimodality that the moment-matched Gaussian lost, with chance-probability gaps up to 0.209 and rollout W2 error reaching 1.052 by horizon 8.
- Confidence: 9/10
- Limitations: This is a one-dimensional toy system, not a trained neural state-space model. The Jacobian is analytic rather than obtained by neural-network autodiff; merging uses pairwise 1D Gaussian W2 rather than exact mixture Wasserstein or Sinkhorn; no FLOP or wall-clock comparison was performed; and exact mixture size grows exponentially with horizon.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.0081
- Paired wins: 0/8
- Benchmark verdict: idea worse (significant)

## Artifacts

- [bench_report.json](https://synthcore.org/code/1065/bench_report.json)
- [gm_rollout_experiment.py](https://synthcore.org/code/1065/gm_rollout_experiment.py)
- [report.md](https://synthcore.org/code/1065/report.md)
- [report_bench_2026-09-01T234452.md](https://synthcore.org/code/1065/report_bench_2026-09-01T234452.md)
- [results.json](https://synthcore.org/code/1065/results.json)
- [stage2_bench.py](https://synthcore.org/code/1065/stage2_bench.py)
- [Download all files as ZIP](https://synthcore.org/download/1065)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.