TEMPORALLY CENTERED SIGREG IMPROVES MULTI-TASK LEWORLDMODEL LEARNING: FROM ANALYSIS TO METHOD

Recent advancements in LeWorldModel (LeWM) have enabled stable, end-to-end world-model learning from pixels without reconstruction or reward supervision. However, these models struggle when scaled to multi-task settings. This paper investigates why the standard regularization technique, the Sketched Isotropic Gaussian Regularizer (SIGReg), fails to maintain effective latent representations when multiple tasks are trained jointly, and proposes a structural fix to enable scalable multi-task learning.

The authors identify that marginal Gaussianization—the process of forcing the entire latent distribution toward an isotropic Gaussian—exerts a contractive pressure that compresses the separation between task-dependent latent clusters. This leads to representation aliasing, where distinct states or tasks become indistinguishable in the latent space. To resolve this, the authors introduce Temporally Centered LeWorldModel (TC-LeWM). Instead of regularizing the full latent marginal, TC-LeWM decomposes the latent into a local temporal mean and a residual, applying SIGReg only to the temporally centered residual. This approach removes the direct pressure on cluster separation while still preventing representation collapse.

TC-LeWM demonstrates significant improvements over the original LeWM on the LIBERO benchmark. In suite-wise 10-task training, the average success rate increased from 53.2% to 73.6%. Furthermore, when scaled to a unified 40-task training regime, TC-LeWM maintained a 73.5% success rate, whereas the original LeWM performance dropped to 44.4%. The authors show that their method produces representations that are more robust to visual perturbations and better organized for downstream behavior-cloning policies, effectively bridging the performance gap between scratch-trained models and large-scale pretrained baselines.