✗ Failed on benchmark 2026

Wasserstein-Controlled Gaussian-Mixture Rollouts

Implementation & benchmark of arXiv:2608.29272 — Stochastic Nonlinear Model Predictive Control with Gaussian Mixture Uncertainty Propagation

Usefulness8/10

Difficulty6/10

Novelty7/10

Dynamics Stability Accuracy Sample efficiency World model Ssm Rl Training dynamics Probability Optimal transport Stochastic processes Statistics Dynamical systems

Source paper: Stochastic Nonlinear Model Predictive Control with Gaussian Mixture Uncertainty Propagation arXiv:2608.29272 ⓘ · analyzed Sep 1, 2026

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

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.

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.$$

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.

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

Failed on benchmark

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗

Stage 2 · Practical benchmark result: Significantly worse ✗

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 9/10

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.

Agent confidence: 9/10

Baseline: Moment-matched single Gaussian: empirical W2 error versus the unmerged mixture was 0.473 at horizon 1, 0.476 at horizon 4, 0.654 at horizon 6, and 1.052 at horizon 8; at d=0.9 its chance estimate was 0.7092.
Idea: The mixture separation matched 2d for d={0.1,0.3,0.6,0.9,1.2}. Below epsilon=1.8 it retained 2 components with empirical W2 error 0.024-0.084; at epsilon>=1.8 it merged to 1 component with error about 0.472. At d=0.9 its chance estimate was 0.5001.

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.

How to run: python3 gm_rollout_experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.29272) Sep 2, 2026

Significantly WORSE than the tuned baseline (+293.3%, p=0.0081, wins 0 / 8 paired seeds).

Benchmark: Dynamics forecasting
Model: rnn_small
Paired seeds: 8
Baseline mean: 0.0002
Idea mean: 0.0007
Effect (Δ): 0.0005 (+293.3%; negative = idea better)
Wins: 0 / 8 paired seeds
p-value: 0.0081 (permutation test, 20 000 shuffles)
Smallest detectable effect: n/a
Mechanism: Not confirmed ✗
Practical effect: Hurts
Baseline tuning: swept over 3 configs

Limitations:

The fixed harness exposes a scalar future-angle target rather than a full long-horizon probabilistic rollout, so the implementation used a learned two-component scalar Gaussian mixture head. It did not implement explicit Jacobian covariance propagation, Wasserstein/Sinkhorn merging, or bounded-component rollout. The run used 800 training examples, 300 test examples, and 15 epochs; larger-scale and longer-horizon tests were not performed.

How to run: python3 stage2_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Sep 1, 2026 latest Failed ✗

Agent confidence: 10/10

Implemented a matched dynamics benchmark using the controlled-pendulum track and the same GRU backbone: deterministic MSE baseline versus a two-component Gaussian-mixture NLL system scored by test MSE. Across 8 paired seeds, the mixture was significantly worse (delta_mean=+0.000501372, permutation p=0.0081), and its trained-model signature did not confirm improved chance calibration. The idea therefore did not transfer as a win on this benchmark.

Baseline: Tuned rnn_small, lr=0.01, 15 epochs: mean test MSE 0.000170967; per-seed [0.000297710, 0.000308663, 0.000086705, 0.000119690, 0.000060004, 0.000177841, 0.000047607, 0.000269516].
Idea: Two-component Gaussian-mixture GRU, best shared-grid lr=0.001, 15 epochs: mean test MSE 0.000672339; per-seed [0.000751154, 0.000564431, 0.001097447, 0.000780845, 0.000161983, 0.000558312, 0.000365519, 0.001099020]. Paired delta +0.000501372; p=0.0081; wins 0/8.

Benchmark result