Technical Summary: QQWorld

Problem Statement

Latent world models enable efficient planning by predicting future states within a compact representation space. A critical design choice in recent architectures, such as LeWorldModel (LeWM), is to regularize the marginal distribution of learned latents toward an isotropic Gaussian. This is typically achieved using the Epps–Pulley (EP) test as a differentiable penalty to enforce normality.

However, the authors observe that despite EP regularization, LeWM often learns latent distributions with pronounced heavy tails. This is problematic because extreme latent values can push learned dynamics into poorly represented regions, amplifying errors during multi-step rollouts and increasing the mismatch with the target Gaussian prior.

The paper identifies a fundamental limitation in the optimization geometry of the EP test:

1. Vanishing Gradients: The EP statistic is equivalent to a squared Maximum Mean Discrepancy (MMD) with a Gaussian kernel. The authors demonstrate that the corrective gradients induced by this objective decay rapidly (super-exponentially) for samples in the tail (outliers).
2. Ineffective Control: Once a latent coordinate moves beyond the interaction scale of the Gaussian kernel, the EP regularizer provides negligible restoring force. Consequently, heavy-tailed deviations are insufficiently controlled, allowing them to persist and grow even when the model is explicitly penalized for non-normality.

Methodology

To address the vanishing gradient issue, the authors propose QQWorld, which replaces the EP regularizer with a Quantile–Quantile (QQ) matching objective.

Core Mechanism: QQ Regularization

Instead of matching moments or using kernel-based discrepancies, QQWorld directly aligns projected latent samples with rank-matched Gaussian quantiles.

• Objective: For a batch of projected samples $X = \{x_n\}$, the loss is defined as the sum of squared differences between the order statistics of the samples ($\hat{x}_n$) and the corresponding theoretical Gaussian quantiles ($q_n = \Phi^{-1}(\frac{n-0.5}{N})$):
  $$L_{QQ}(X) = \sum_{n=1}^N (\hat{x}_n - q_n)^2$$
• Gradient Behavior: Unlike the EP gradient, which vanishes for extreme values, the gradient of the QQ loss with respect to a sample $x_n$ is $2(x_n - q_{\rho(n)})$, where $\rho(n)$ is the rank. This gradient magnitude grows linearly with the deviation from the target quantile. This ensures that tail samples receive strong, symmetric corrective forces, effectively suppressing heavy tails.
• Theoretical Guarantee: The authors prove that driving the QQ loss to zero also drives the EP loss to zero (one-way control), but the converse is not true. A small EP loss can coexist with arbitrarily large tail deviations, whereas QQ matching enforces stricter distributional alignment.

Cross-Batch QQ

To address the memory constraints associated with large batch sizes (which improve the accuracy of rank estimation), the authors introduce Cross-Batch QQ.

• Strategy: This method enlarges the effective ranking pool by including detached features from previous $K$ iterations into a FIFO queue.
• Mechanism: Ranks are computed over the pooled set of size $M = N(K+1)$, but gradients are backpropagated only through the current batch of size $N$.
• Bias-Variance Trade-off: The authors characterize a trade-off where increasing the queue length $K$ reduces the variance of the rank estimation (beneficial for small batches) but introduces a "representation-staleness" bias if the latent distribution shifts significantly over time.

Key Contributions

1. Identification of EP Limitation: The paper formally analyzes and demonstrates that the EP regularizer's corrective force vanishes for tail samples, explaining why LeWM fails to suppress heavy tails despite explicit normality penalties.
2. Proposal of QQWorld: A drop-in replacement for LeWM that utilizes a quantile–quantile matching objective. It maintains effective corrective gradients for tail deviations without introducing new hyperparameters.
3. Memory-Efficient Training: The development of Cross-Batch QQ, which allows for accurate quantile estimation using large effective batch sizes while keeping GPU memory overhead low by decoupling the ranking pool size from the backpropagation batch size.
4. Theoretical Analysis: Characterization of the bias-variance trade-off in Cross-Batch QQ and proof of the one-way control relationship between QQ and EP objectives.

Experimental Results

The authors evaluated QQWorld across four control environments: Two-Room, PushT, Reacher, and OGBench-Cube.

• Planning Performance: QQWorld achieved an average planning success rate of 85.08%, outperforming the LeWM baseline (79.75%) by 5.33 percentage points. It consistently surpassed LeWM and other recent variants (Sub-JEPA, SD-JEPA, SMWM) across all four environments.
• Distributional Alignment:
  • QQ RMSE: Reduced the mean error between empirical and Gaussian quantiles from 0.157 (LeWM) to 0.121 (QQWorld).
  • Tail Probability: Reduced the probability of samples exceeding the 90th percentile of the chi-squared distribution (a measure of radial tail heaviness) from 0.315 to 0.123, approaching the Gaussian reference of 0.10.
  • Statistical Tests: QQWorld reduced the Kolmogorov–Smirnov (KS) statistic by 15.8% and the EP statistic by 31.4% compared to LeWM, even though LeWM was trained directly with the EP objective.
• Physical State Probing: On the PushT environment, QQWorld matched or slightly outperformed LeWM in decoding physical states (agent location, block location, block angle), indicating that the stronger regularization did not degrade the encoding of task-relevant information.
• Cross-Batch Efficiency: Using a small batch size ($N=32$) with a queue length of 3 ($K+1=3$), QQWorld achieved an 83.50% success rate, outperforming LeWM trained with a batch size of 128 (79.75%), while reducing GPU memory usage by 73%.

Significance and Claims

The paper claims that the effectiveness of a statistical test for measuring distributional discrepancy does not guarantee its efficacy as a training objective. Specifically, the authors argue that:

• Gradient Quality Matters: Beyond merely distinguishing distributions, a regularization objective must provide informative and well-behaved gradients throughout the optimization process. The EP test fails this criterion for tail samples, whereas the QQ objective succeeds.
• Tail Control is Critical: Effective suppression of heavy tails in latent spaces is essential for stable multi-step planning in world models.
• Scalability: The Cross-Batch QQ strategy offers a scalable mechanism for training large world models where increasing physical batch sizes is prohibitively expensive in terms of memory.

The work concludes that QQWorld provides a simple, hyperparameter-free improvement over existing latent world models by directly addressing the optimization geometry of distributional regularization.