Temporally Centered SIGReg Improves Multi-Task LeWorldModel Learning: From Analysis to Method

Overview

LeWorldModel (LeWM) enables end-to-end joint-embedding world-model learning directly from pixels. It uses the Sketched Isotropic Gaussian Regularizer (SIGReg) to prevent representation collapse by encouraging the latent marginal to follow an isotropic Gaussian.

We find that this marginal target does not extend reliably to multi-task training. Marginal Gaussianization compresses the separation between task-dependent latent clusters relative to within-cluster variation, which can introduce representation aliasing across tasks and states despite global non-collapse.

We introduce Temporally-Centered LeWorldModel (TC-LeWM). TC-LeWM applies SIGReg to temporally centered residuals rather than to the full latent marginal. This removes direct regularization pressure on low-frequency cluster centers while retaining the anti-collapse effect of SIGReg.

Across the four LIBERO suites, TC-LeWM raises the suite-wise 10-task average success rate from 53.2% to 73.6%. Under unified 40-task training, it maintains 73.5% average success, compared with 44.4% for Raw LeWM.

Multi-task degradation in Raw LeWM

We compare independent single-task training with joint training on the same ten LIBERO-Long tasks. Raw LeWM achieves 40.0% average downstream success when a separate pipeline is trained for each task, but only 29.9% when all ten tasks share one encoder, predictor, and task-conditioned policy. Eight of the ten tasks deteriorate under joint training.

The latent representation remains statistically non-collapsed, but global non-collapse alone does not ensure that the representation preserves the state distinctions required by a shared downstream behavior-cloning policy.

Why marginal SIGReg can degrade multi-task learning

We decompose each latent into a local temporal mean and a complementary short-timescale residual, z_t = z̄_t + r_t. For a random one-dimensional projection, linearity gives x = c + ε, where c represents low-frequency task- and context-dependent structure and ε represents complementary short-timescale variation.

We model c as a balanced, homoscedastic K-component Gaussian mixture and vary the normalized global component-center spread, ρ = std_k[μ_k] / σ. We evaluate K ∈ {2, 5, 10, 20} over ρ ∈ {0, 0.1, …, 4.0}, using 500 randomized center configurations for each mixture size.

Explore how the projected mixture and the EP response change with the normalized spread ρ:

As ρ approaches 1 or below, the global center spread becomes comparable to the within-component standard deviation, making component overlap and representation aliasing likely. On LIBERO-Long, Raw LeWM yields ρ = 0.74, whereas TC-LeWM yields ρ = 2.54. This contrast is consistent with the latent geometries examined below.

For a task-conditioned behavior-cloning policy π(a_t | z_t, i), the task identity can resolve ambiguity attributable solely to the task label, but it cannot recover within-task state distinctions already aliased in z_t. Similar latent codes may therefore correspond to incompatible actions.

TC-LeWM: from marginal to residual

TC-LeWM retains the LeWM encoder and action-conditioned latent predictor. For each latent sequence, it computes a local temporal mean z̄_t and the centered residual r_t = z_t − z̄_t. Raw LeWM and TC-LeWM differ only in the target of SIGReg:

No direct pressure on low-frequency separation. In the projected mixture model, changing ρ scales the component centers of c while leaving the residual distribution ε unchanged. The residual EP objective is therefore invariant to ρ under this construction. TC-SIGReg avoids the center-contraction mechanism of marginal SIGReg without explicitly forcing task-dependent clusters apart.

Anti-collapse is retained. If the latent becomes constant within a temporal window, then r_t ≡ 0 and every residual projection becomes a degenerate point mass, which EP penalizes. The prediction and residual-regularization objectives together favor non-degenerate, predictable temporal variation without directly Gaussianizing low-frequency multi-task structure.

Quantitative results

Closed-loop success rates (%) under suite-wise 10-task training. Raw LeWM and TC-LeWM are reported as mean±std over three full-pipeline training seeds. Diffusion Policy, Octo, and OpenVLA are published mean results reported in Kim et al. (2025) and are included as contextual references.

Under suite-wise 10-task training, TC-LeWM improves success over Raw LeWM across all four suites, raising the average from 53.2% to 73.6%. The largest improvements occur on LIBERO-Object and LIBERO-Long: +41.4 and +21.9 percentage points.

As task scale increases, Raw LeWM degrades while TC-LeWM remains stable. On LIBERO-Long, Raw LeWM changes from 40.0% under single-task training to 29.9% under 10-task joint training, whereas TC-LeWM changes from 49.6% to 51.8%. From suite-wise 10-task training to unified 40-task training, Raw LeWM decreases from 53.2% to 44.4%, while TC-LeWM remains nearly unchanged at 73.6% to 73.5%.

Closed-loop success rates (%) under unified 40-task training. All four LIBERO suites share one encoder, predictor, and task-conditioned policy. Results are mean±std over three full-pipeline training seeds; the final column reports the change from suite-wise 10-task training.

Closed-loop rollouts

Examples from the suite-wise 10-task setting, comparing Raw LeWM and TC-LeWM on the same tasks under the same evaluation protocol.

Example 1 Example 2 Example 3

4× speed, looping

Raw LeWM 47.6% success (450 rollouts)

FAIL

SUCCESS

SUCCESS

FAIL

TC-LeWM (ours) 69.1% success (450 rollouts)

FAIL

SUCCESS

SUCCESS

SUCCESS

Within a column both models are evaluated from the same initial scene, so the clip above and the clip below it are directly comparable.

Latent structure and perturbation robustness

We examine how changing the SIGReg target affects latent-space organization, short-timescale variation, and robustness to visual perturbations.

@misc{liu2026tcsigreg,
  title         = {Temporally Centered SIGReg Improves Multi-Task LeWorldModel Learning: From Analysis to Method},
  author        = {Chang Liu and Fei Suo and Yanzhou Jin and Yusuke Iwasawa and Yutaka Matsuo and Yaonan Zhu},
  year          = {2026},
  eprint        = {2607.26924},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2607.26924}
}