PaperlayerView exact revision on arXiv

Source-linked AI summary

Spectral-Target Physical Latent Structuring for JEPA-Style World Models

Penghao Zhu, Salvatore Penachio, Kaustav Mukherjee, Aneesh Jonelagadda

arXiv:2609.04264v3cs.LG

TL;DR

Preventing collapse in a latent world model does not ensure that its representation preserves physical information needed for planning. This paper adds training-time Fourier auxiliary supervision on task-relevant object positions or boxes; experiments report increased planning success, particularly where LeWM is physically lazy and in data-limited regimes.

  • Problem

    Regularization can prevent latent collapse without ensuring that the representation preserves physical variables needed for planning.

    Source: §1 · pp. 1–2

  • Method

    A lightweight Fourier auxiliary head uses task-relevant object positions or bounding boxes to supervise encoded and predicted latent states during training.

    Source: §1 · pp. 2–3

  • Results

    Planning success improves across LeWM environments, with the n-ball results showing gains where the baseline has physical representation laziness.

    Source: §5.2 · p. 8, §6 · p. 9

  • Takeaways & Limitations

    The results support preserving task-relevant physical state in addition to avoiding representation collapse, with particular value in data-limited regimes.

    Source: §1 · p. 3

Abstract

from arXiv · show

Latent world models have become increasingly popular as a method to predict and plan in latent space rather than pixel space. Recent architectures, such as LeWorldModel (LeWM), jointly train the encoder and predictor using regularization techniques like SIGReg to prevent representation collapse. Even with such regularization preventing representation collapse, we identify a new world model failure mode of physical representation laziness, particularly noted in highly dynamic environments. For these lazy cases, the learned latent states do not collapse but nonetheless fail to represent key physical properties, causing ubiquitous downstream planning failure. To resolve this issue, we propose training-time auxiliary supervision with a lightweight "Fourier auxiliary head", which enforces physically-informed structuring of the latent space with no additional inference-time cost and can be generalized to any environment. Experimentally, we show that the auxiliary head substantially improves planning success rates in dynamic environments where the baseline LeWM exhibits physical representation laziness. It also leads to modest improvements in other environments, even when the baseline does not exhibit physical representation laziness. We further observe superior planning performance being accompanied by higher latent space correlations with key physical properties, indicating both the ability of our method to physically structure latent states and the potential planning-side benefit to the learned representation being physically structured. We also see in low-data regimes, auxiliary supervision is particularly impactful in increasing success rate. These findings support the use of our Fourier auxiliary head method to improve both overall success rate and data efficiency, while avoiding representation laziness in latent world models.

1 INTRODUCTION

JEPA-style world models predict in latent space, but preventing representation collapse does not ensure that the encoder retains physical information needed for planning. The paper introduces a lightweight Fourier auxiliary head to structure latents around task-relevant objects, improving physical grounding and planning without adding inference-time cost.

1 INTRODUCTION

  • The Fourier auxiliary head trains a small MLP to recover spectral features of task-relevant object positions or bounding boxes from encoded states.The same supervision can be applied to predicted latent states to encourage preservation of physical information over time.Source: §1 · pp. 2–3
  • SIGReg encourages latent diversity, whereas the auxiliary objective asks whether the representation contains the physical information required by the task.The objectives address different failure modes: global collapse versus treating physically important observations as effectively identical.Source: §1 · pp. 1–2
  • The head is primarily for training and can be removed afterward without changing the encoder–predictor architecture.When explicit physical rules exist, it can instead be retained to translate predicted positions into planning costs; this paper focuses on the indirect mechanism.Source: §1 · p. 3
  • Physical grounding improves encoder representations, reduces latent prediction error, and increases downstream planning success, with particular value in data-limited regimes.The authors conclude that avoiding collapse is necessary but not sufficient: a JEPA must also preserve physical state relevant to decisions.Source: §1 · p. 3
  • Physical representation laziness occurs when a non-collapsed encoder omits physical information needed for prediction and planning.Unlike a collapsed encoder, a lazy encoder produces different representations but may ignore objects or geometric relations relevant to planning.Source: §1 · pp. 2–3

2 RELATED WORK

Related work spans reconstruction-based and decoder-free latent world models, as well as joint-embedding predictive architectures. This paper builds on end-to-end latent prediction while addressing physical variables required for planning, complementing other extensions and forms of geometric or auxiliary supervision.

World models

  • Reconstruction-based world models learn compact latent dynamics with pixel reconstruction or prediction, while TD-MPC2 and DINO-WM use decoder-free alternatives.DreamerV3 scales the reconstruction-based approach across diverse domains; TD-MPC2 uses task-specific consistency, reward, and value objectives, and DINO-WM models features from a fixed pretrained DINOv2 encoder.Source: §2 · p. 3

Joint-embedding predictive models

  • Joint-embedding predictive models predict target representations rather than reconstructing pixels, and LeWM jointly trains an encoder and predictor with SIGReg.The paper studies a distinct failure mode: non-collapsed representations can still omit physical variables required for planning.Source: §2 · p. 3

Extensions of LeWM

  • Other LeWM extensions target planning speed, hierarchical subgoals, online adaptation, latent regularization, or planning-useful latent distances.Fast-LeWM, Hi-LeWM, and AdaJEPA address planning in complementary ways; RC-aux and Temporal-Distance JEPA introduce auxiliary objectives [Li et al., 2026; Bai & Xiong, 2026].Source: §2 · pp. 3–4
  • Depth regularization and proprioceptive grounding provide related geometric supervision, whereas this work grounds task-relevant object positions or boxes in encoded and predicted states.The paper also considers limited data, visually indistinguishable objects, and planning with explicit physical rules.Source: §2 · p. 4

3 BEYOND REPRESENTATION COLLAPSE: A CASE STUDY OF ENCODER LAZINESS IN JEPA WORLD MODELS

In a two-ball task, planning requires representing both a controlled ball and an autonomously moving environmental ball. The case study shows that SIGReg can prevent collapse while the latent geometry fails to preserve the controlled ball’s position, motivating physical grounding.

3 BEYOND REPRESENTATION COLLAPSE: A CASE STUDY OF ENCODER LAZINESS IN JEPA WORLD MODELS

  • ρenv = 0.4006 and ρctrl = −0.0008: latent distance tracks environmental-ball position but not controlled-ball position despite SIGReg.Figure 2 plots latent distance against each ball’s position difference; the task requires moving the controlled ball to a target without colliding with the environmental ball.Source: Figure 2 · p. 4, §3 · p. 4
  • The encoder is non-collapsed but physically lazy: it captures only one ball’s position although planning requires both.The controlled ball moves in response to actions, while the environmental ball moves autonomously and bounces off walls.Source: §3 · pp. 4–5
  • SIGReg prevents global collapse but does not determine which physical variables the encoder preserves, leaving the latent space unreliable as a planning objective.The case study measures latent distances between encoded observation pairs and compares them with each ball’s position differences.Source: §3 · pp. 4–5

4 METHODS

The methods add geometric auxiliary supervision to JEPA-style world models by converting object boxes into permutation-aware Fourier targets, shaping latent representations during training while planning in latent space.

4 METHODS

  • Bounding-box states for exchangeable objects pass through a fixed Fourier feature map and are summed into the spectral target; auxiliary heads are removed at inference.Source: Figure 3 · p. 5, §4 · p. 5

4.1 JEPA-STYLE WORLD MODELS

  • LeWM uses MSE prediction loss and SIGReg to prevent collapse, but these regularizers do not ensure that latents preserve planning-relevant information.Source: §4.1 · pp. 5–6

4.2 FOURIER-SYMMETRIZED BOUNDING BOX AUXILIARY HEAD

  • The method symmetrically aggregates bounding-box states with a Fourier sum, producing a permutation-invariant spectral representation for visually similar or identical objects.Source: §4.2 · p. 6

4.3 ARCHITECTURE OF THE AUXILIARY HEAD

  • A two-layer MLP predicts spectral targets, and auxiliary losses train the encoder and predictor; the head stays shallow and is removed at inference.Source: §4.3 · p. 7

4.4 TRAJECTORY PLANNING

  • CEM optimizes action sequences by rolling out predicted latents; for tasks with explicit rules, predicted geometry can also augment the planning objective.Source: §4.4 · p. 7

5 RESULTS

The auxiliary head improves planning success across environments, with particularly strong gains in n-ball tasks and sustained benefits in lower-data Push-T settings. Latent-state correlations generally improve, though gains vary by environment and data availability.

5.1 EVALUATION SETUP

  • Evaluation covers 1-, 3-, and 6-ball tasks, Push-T, Two Room, and OGBench Cube, with planning success averaged across five seeds.The environments use different training-set sizes; latent-distance correlations assess how well latent distances track object-state differences.Source: §5.1 · p. 7, §5.1.1 · p. 8

5.2 PLANNING SUCCESS

  • Planning success improves across environments, with the n-ball results highlighting LeWM's physical representation laziness and the auxiliary target's benefits.The reported n-ball gains coincide with better organization of latent states, although environment-ball positional understanding is slightly worse for n = 1 and n = 3.Source: Table 1 · p. 8, §5.2 · p. 8, Table 2 · p. 8
  • Other environments show modest but consistent planning gains; Spearman correlations improve for Push-T and Two Room, while Reacher and OG-Bench Cube latent spaces collapse despite improved success.Strong correlations for Reacher and OG-Bench Cube at lower data amounts, even 500k, are also reported.Source: Table 3 · p. 8, §5.2 · p. 9

5.3 DATA EFFICIENCY

  • The Push-T auxiliary-head advantage persists at smaller training-data sizes, but at sufficiently low data levels the auxiliary target cannot improve baseline results because of underdetermination.The caption reports improvements over LeWorldModel in lower data regimes.Source: §5.3 · p. 9, Figure 4 · p. 9

6 CONCLUSION

The paper introduces Fourier auxiliary heads to structure JEPA-style world-model latents around object properties, improving planning without additional inference cost. It reports the largest planning gains in n-ball tasks and notes that real-world environments and other JEPA baselines remain untested.

6 CONCLUSION

  • Planning success rises by 91.3%, 78.5%, and 65.3% (H=4) for 1, 3, and 6 environmental balls, respectively, over LeWM.The LeWM baseline exhibits encoder laziness in the n-ball environment.Source: §6 · p. 9
  • Fourier auxiliary heads use training-time supervision derived from bounding boxes to organize latents for object-centric understanding without additional inference cost.The spectral target is described as generalizable to 2D and 3D environments.Source: §6 · p. 9
  • The method also consistently improves task performance in tested common environments at full and low-data settings, with improved physical-property encoding in some benchmarks.The conclusion reports near-zero control-ball correlations for LeWM and adequate correspondence between latent states and ball position with the method.Source: §6 · p. 9
  • The method has not been tested in real-life environments, and the experiments use LeWM as the only baseline.The conclusion leaves combinations with additional supervision, regularization, and architectures for future work.Source: §6 · p. 9

A.1 TRAINING AND EVALUATION PIPELINE CONFIGURATION

The pipeline uses two auxiliary-head architecture variants, distinguishing n-ball tasks from the remaining environments.

A.1 TRAINING AND EVALUATION PIPELINE CONFIGURATION

  • The auxiliary-head architecture has two variants.One is used for the n-ball environment and another for the remaining environments.Source: §A.1 · p. 13
  • The n-ball environment uses its own auxiliary-head architecture variant.Source: §A.1 · p. 13
  • For the remaining environments, the architecture is kept as close as possible to the original LeWM implementation.Source: §A.1 · p. 13

A.1.1 N-BALLS ARCHITECTURE

The n-ball setup specifies image inputs, encoder and predictor designs, auxiliary-head targets, loss weights, and a fixed-horizon planning protocol.

A.1.1 N-BALLS ARCHITECTURE

  • Rendered RGB frames are 64 × 64, normalized to [−1, 1]; training uses batch size 64 for 200 epochs with four sub-trajectories of four observations and actions each.No frameskip is used.Source: §A.1.1 · p. 13
  • A custom vision transformer uses patch size 8, width 64, four blocks, four attention heads, 256-D FFNs, and a CLS token projected to a 128-D latent.Source: §A.1.1 · p. 13
  • The action-conditioned causal-transformer predictor takes three history frames, uses four blocks and four attention heads, and projects its output to 128 dimensions.It uses AdaLN for action conditioning and causal masking.Source: §A.1.1 · p. 13
  • Two-layer auxiliary MLPs map the 128-D latent through a 64-D feature to an 18-D spectral target; λAux is 0.1 and λSIGReg is 0.09.Source: §A.1.1 · p. 13
  • Planning uses CEM with 100 candidate sequences, 10 elites, and 20 iterations; success means the control ball reaches within 1 pixel of the target.The action norm is clipped at 3.0, evaluation uses 200 episodes per seed, and horizons are 4, 6, or 8 without replanning.Source: §A.1.1 · p. 13

A.1.2 LEWM ARCHITECTURE

The LeWM-style setup uses ViT-based encoder and predictor models, a lightweight spectral auxiliary head, and CEM planning. Models are trained on rendered environment frames with a fixed observation-and-action sequence format.

A.1.2 LEWM ARCHITECTURE

  • Rendered RGB frames are trained at 224 × 224 resolution for 10 epochs, using batch size 128 and four sub-trajectories of four observations and four five-action blocks.A frameskip of 5 is used, so five actions are predicted for every frame.Source: §A.1.2 · p. 13
  • The encoder is ViT-Tiny with patch size 14 and a 192-D CLS-token latent; the predictor is ViT-Small with causal masking and learned position embeddings.The predictor accommodates a history length of three.Source: §A.1.2 · p. 13
  • Two-layer MLP auxiliary heads map the 192-D latent through a 64-D feature to an 18-D spectral target, with λAux = 0.1 and λSIGReg = 0.09.Source: §A.1.2 · p. 13
  • CEM planning samples 300 action sequences, keeps 30 elites, and runs 30 iterations; MPC replans every 25 actions under a 50-step budget.The action-sampling variance scale is 1.0; environments can be solved in 25 steps.Source: §A.1.2 · pp. 13–14

A.1.3 TARGET CREATION

Targets are ground-truth simulator bounding boxes, selected according to which objects matter in each environment. Figure 5 shows sample targets for the tested environments.

A.1.3 TARGET CREATION

  • All balls are targets in n-ball environments, while only dynamic objects are targeted elsewhere; the boxes come from simulator ground truth.Figure 5 shows example targets for each environment.Source: §A.1.3 · p. 14, Figure 5 · p. 14

A.2 THEORETICAL EXAMINATION OF VISREG + AUX HEAD

The appendices examine how auxiliary supervision interacts with regularization, planning constraints, target choices, architecture, object matching, depth, noisy labels, and physical probing. Across these settings, the results indicate benefits from physical supervision while exposing differences among target representations and configurations.

A.2 THEORETICAL EXAMINATION OF VISREG + AUX HEAD

  • The auxiliary method performs worse with VISReg than with SIGReg, and the proposed explanation is that weak SIGReg corrections may let task-induced spectral anisotropy persist.The authors contrast VISReg’s isotropy objective with their method’s target of spectral anisotropy reflecting spatial structure.Source: Figure 6 · p. 15, §A.2 · pp. 15–16

A.3 AUXILIARY HEAD IN CEM PLANNING

  • Auxiliary-informed CEM raises success rates across all n-ball experiments, especially with more balls, while improving collision avoidance at some cost to target distance for higher n.It uses auxiliary predictions from imagined latents to penalize collisions, without providing CEM ground-truth ball positions during planning.Source: §A.3 · p. 17, Table 4 · p. 18

A.4 TARGET ABLATIONS

  • Bounding boxes yield the highest success rates in Push-T and 3-Balls, while Compact13 gives the highest success rate for OG-Bench Cube.The 2D box-size depth proxy structures the latent less effectively than explicit 3D supervision; such privileged information may not be available to generalize.Source: Table 5 · p. 19, §A.4 · p. 19

A.5 AUX HEAD ARCHITECTURE ABLATIONS

  • Auxiliary gradients into the predictor improve planning over a single auxiliary head, though the single-head variant still improves success over baseline LeWM.The final method shares an auxiliary head across decoded and predicted latents.Source: Table 6 · p. 20, §A.5 · p. 20

A.6 OBJECT MATCHING ABLATIONS

  • Fourier symmetrization consistently outperforms no matching; frame-wise Hungarian matching varies by task, underperforming no matching on OG-Bench Cube and collapsing on 3-Balls.The ablation compares Fourier symmetrization with no matching and frame-wise or sequence-wise Hungarian matching.Source: Table 7 · p. 20, §A.6 · p. 20

A.7 DEPTH ABLATIONS

  • Adding 25th- and 75th-percentile depth to boxes marginally improves performance but remains less useful than Compact13 object centroids on OG-Bench Cube.Source: Table 8 · p. 21, §A.7 · p. 21

A.8 NOISY BOUNDING BOX ABLATION; A.9 ZERO-SHOT OBJECT DETECTION PIPELINE FOR BOUNDING BOX DATA

  • Noisy boxes retain above-baseline success and prevent 3-Balls collapse; zero-shot detections closely resemble simulator labels, demonstrating feasibility without manual frame-level annotation or simulator labels.The noise experiment jitters each box edge by up to 1% of image width.Source: §A.8 · p. 21, Table 9 · p. 21, §A.9 · p. 22

A.10 LINEAR AND MLP PROBING

  • Probes improve for most physical quantities across environments, with Push-T and 6-ball benefiting most; Two-Room has poorer linear but improved MLP probe performance.The authors interpret the Two-Room pattern as more entangled but improved positional encoding.Source: Table 10 · p. 23, §A.10 · p. 23
Loading 2609.04264v3…