QQWorld: Quantile-Quantile Matching for World Model Regularization

Summary

Abstract Summary

Latent world models enable efficient planning by predicting future states in a compact representation space, but their performance depends critically on the quality of the learned latent distribution. LeWorldModel (LeWM) regularizes its latents toward an isotropic Gaussian using the Epps-Pulley (EP) objective. We show that the corrective gradients of EP rapidly vanish for isolated tail samples, leaving heavy-tailed deviations insufficiently controlled. To address this limitation, we propose QQWorld, which replaces EP with a quantile-quantile matching objective that directly aligns projected latent samples with rank-matched Gaussian quantiles, thereby maintaining effective corrective gradients in the tails. We further develop cross-batch QQ, which enlarges the effective ranking pool using detached samples from previous batches, and characterize its bias-variance trade-off. Across four control environments, QQWorld effectively improves the average planning success rate of LeWM, while consistently yielding better Gaussian alignment and thinner latent tails.

Key Points

Latent world models enable efficient planning by predicting future states in a compact representation space, but their performance depends critically on the quality of the learned latent distribution.
LeWorldModel (LeWM) regularizes its latents toward an isotropic Gaussian using the Epps-Pulley (EP) objective.
We show that the corrective gradients of EP rapidly vanish for isolated tail samples, leaving heavy-tailed deviations insufficiently controlled.
To address this limitation, we propose QQWorld, which replaces EP with a quantile-quantile matching objective that directly aligns projected latent samples with rank-matched Gaussian quantiles, thereby maintaining effective corrective gradients in the tails.
We further develop cross-batch QQ, which enlarges the effective ranking pool using detached samples from previous batches, and characterize its bias-variance trade-off.

Method

Let $z_{1},\dots,z_{N}\in\mathbb{R}^{d}$ denote the latent embeddings of a training batch. Similar to Balestriero and LeCun (2025) , normality is enforced along random one-dimensional projections.

For a direction $u\sim\operatorname{Uniform}(\mathbb{S}^{d-1})$ where $\mathbb{S}^{d-1}$ denotes the unit sphere in $\mathbb{R}^{d}$, we form the projected samples $x_{n}=\langle u,z_{n}\rangle$ and penalize their departure from $\mathcal{N}(0,1)$.

Method — QQWorld: Quantile-Quantile Matching for World Model Regularization — Figure 3 : Visualization of the learned latent representations on Two-Room. We train a linear probe to predict the two-dimensional agent location from a latent representation. The figure visualizes two representative test episodes. For each episode, the left subpanel shows the ground-truth trajectory, while the right subpanel magnifies the region enclosed by the dashed box. Gray curves denote the ground-truth trajectories, and colored curves denote the trajectories decoded from the learned latents. The color along each predicted trajectory represents the cumulative RMSE up to the corresponding time step. The value below each example reports its episode RMSE, while the leftmost summary column reports the RMSE computed over all test episodes.

Experiments

We follow LeWM ( Maes et al., 2026 ) for dataset preprocessing, model training, and CEM-based goal-conditioned evaluation. Experiments are conducted on the same offline datasets, including Two-Room ( Sobal et al., 2025 ) , PushT ( Zhou et al., 2025 ) , Reacher ( Tassa et al., 2018 ) , and OGBench-Cube ( Park et al., 2025 ) . QQWorld differs from LeWM only by replacing the EP regularizer with the proposed QQ objective.

We set the QQ regularization weight to $3.5$ for all environments.

Sources and Demonstrations

Figures are reproduced from Yu et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.