QQWorld: Quantile-Quantile Matching for World Model Regularization

2026-07-30

AI Summary

This work addresses the insufficient regularization of heavy-tailed distributions in the latent space of existing world models, which degrades planning performance. To mitigate this issue, the authors propose a quantile–quantile (QQ) matching–based regularization method that directly aligns latent variable samples with standard Gaussian quantiles, thereby enhancing gradient correction in tail regions. A cross-batch sampling mechanism is further introduced to expand the ranking pool and strengthen regularization efficacy. The proposed approach effectively alleviates gradient vanishing for tail samples, significantly improving the average planning success rate of LeWM across four control environments while achieving superior Gaussian alignment and thinner tails in the latent distribution.

Abstract

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.

Problem

• latent world models
• regularization
• heavy-tailed deviations
• Gaussian alignment
• tail samples

Innovation

• quantile-quantile matching
• world model regularization
• latent distribution alignment
• cross-batch ranking
• heavy-tailed correction