QQWorld Regularizes Latent World Models with Quantile Matching

Introduction

Latent world models make planning more efficient by predicting future states in a compact representation space. But the quality of that space is not only about compression or prediction error. Its distributional shape matters: if latent variables develop heavy tails or deviate from the Gaussian structure assumed by the model, imagined rollouts may become less reliable, which can affect downstream control performance.

QQWorld addresses this issue in the context of LeWorldModel, or LeWM. LeWM regularizes latent variables toward an isotropic Gaussian distribution using the Epps-Pulley objective. The new paper argues that this choice has a weakness: for isolated samples in the distribution tails, the corrective gradients of EP can decay rapidly. As a result, heavy-tailed deviations may persist even while the model is nominally being regularized toward a Gaussian.

Key ideas

• The problem is in the tails. EP can encourage Gaussian-like latents overall, but the paper highlights that samples far from the main mass of the distribution may receive weak corrective pressure. These tail samples are important because world models often rely on repeated rollout, where small distributional flaws can accumulate during planning.

• QQWorld replaces EP with quantile-quantile matching. Instead of using the EP objective, QQWorld projects latent samples, ranks them, and aligns each sample with the Gaussian quantile corresponding to its rank. This makes the regularization more direct: the model is asked to match not just broad statistics, but the shape of the projected distribution.

• Tail correction remains effective. Because rank-matched quantiles explicitly include the extremes, tail samples receive a clearer target. This is the central motivation of QQWorld: it preserves useful corrective gradients for samples that would otherwise be poorly controlled by the previous objective.

• Cross-batch QQ improves the ranking pool. Quantile matching depends on the available samples used for ranking. To make this more stable, the authors introduce cross-batch QQ, which augments the current batch with detached samples from previous batches. The paper also discusses the bias-variance trade-off created by this design.

Why it matters

Across four control environments, QQWorld improves the average planning success rate of LeWM while producing better Gaussian alignment and thinner latent tails. The provided material does not include exact numbers, so the main takeaway is qualitative but important: distributional regularization in latent world models can meaningfully affect planning quality.

This work is especially relevant for reinforcement learning, robotics, and embodied AI settings, where a world model is used to imagine future trajectories before acting. If the latent space contains poorly controlled heavy-tail behavior, imagined futures may drift away from useful dynamics. By directly aligning samples with Gaussian quantiles, QQWorld offers a targeted way to reduce that risk.

The contribution is not a completely new world model architecture. Instead, it refines a crucial training component: how the latent space is regularized. That makes the paper interesting because it shows how a mathematically focused change in the objective can translate into more stable planning behavior.

Source: Hugging Face Daily Papers