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

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 \textit{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.

한 줄 요약

JEPA 기반 월드 모델에서 물리적 정보를 효과적으로 인코딩하기 위해 Fourier auxiliary head를 도입하여 planning 성능과 데이터 효율성을 향상시킨다.

핵심 기여도

**Physical representation laziness** 문제를 식별: SIGReg가 representation collapse를 방지하더라도, 핵심 물리적 정보를 생략할 수 있음.
**Fourier auxiliary head** 제안: task-relevant object의 위치나 bounding box를 기반으로 latent space에 물리적 구조를 강제.
**Planning 성능 개선**: 1, 3, 6개의 ball 환경에서 각각 88.8%, 69.6%, 62.7%의 planning success rate 증가.
**Low-data regime에서 효과**: 추가적인 감독 없이도 데이터 효율성을 향상.

핵심 아이디어

JEPA 기반 월드 모델은 latent space에서 예측과 계획을 수행하지만, 기존의 SIGReg는 representation collapse를 방지하지만, 핵심 물리적 정보를 생략하는 **representation laziness**를 막지 못한다. 이는 downstream planning 성능에 부정적인 영향을 미친다. 이를 해결하기 위해, 저자는 **Fourier auxiliary head**를 제안한다. 이 모듈은 training 시에만 사용되며, task-relevant object의 위치나 bounding box를 기반으로 latent state에서 **spectral features**를 복원하도록 학습한다. 이는 latent space가 물리적 정보를 반영하도록 유도하며, planning 성능을 향상시킨다. 특히, 이 모듈은 **inference 시 제거 가능**하며, 추가적인 계산 비용이 발생하지 않는다.

기술적 접근법

**Fourier auxiliary head**: 2-layer MLP로 구성.
**Spectral target**: object의 위치 또는 bounding box를 기반으로 생성.
**Loss function**:
$ \mathcal{L}_{Aux,t} $: 각 latent state $ z_t $에 대한 auxiliary loss.
$ \lambda_{A1}, \lambda_{A2} $: auxiliary loss를 encoder와 predictor에 통합하는 가중치.
**SIGReg**: latent distribution이 isotropic Gaussian과 유사하도록 강제.
**Training data**: ball, Push-T, Two Room, OGBench Cube 환경에서 각각 100k~730k 데이터 사용.

주요 결과

**Planning success rate**:
1-ball 환경에서 88.8% 증가.
3-ball 환경에서 69.6% 증가.
6-ball 환경에서 62.7% 증가.
**Latent distance correlation**:
LeWM 기반 모델은 control ball의 위치와의 상관계수 $ \rho $가 거의 0.
Fourier auxiliary head를 사용하면 $ \rho $가 상승, latent space가 물리적 정보를 반영함을 확인.
**Low-data regime**: 100k 이하 데이터에서도 성능 개선 효과가 두드러짐.

의의 및 한계

**의의**:
SIGReg가 방지하지 못하는 **representation laziness** 문제를 해결.
물리적 정보를 반영한 latent space를 학습하여 planning 성능을 향상.
추가적인 inference cost 없이도 학습 과정에서 효과적으로 활용 가능.
**한계**:
실제 환경에서의 low-data regime 성능 검증 미흡.
LeWM만 기반으로 실험, 다른 JEPA 기반 모델과의 호환성 검증 필요.
Fourier-symmetrized target의 일반화 가능성은 제시되었으나, 다양한 환경에서의 검증 필요.

실용적 활용

**로봇 제어 및 자율 주행**: 물리적 정보를 반영한 latent space를 통해 planning 성능 향상.
**시뮬레이션 기반 학습**: 데이터가 제한적인 상황에서 효과적인 world model 학습 가능.
**object-centric representation 학습**: 복잡한 환경에서도 물리적 구조를 반영한 인코딩 가능.