Reproduce “UWM-JEPA: Predictive World Models That Imagine in Belief Space” on OpenResearch

About this paper

Santosh Kumar Radha, Oktay Goktas

Researchers introduced UWM-JEPA, a Joint Embedding Predictive Architecture that uses density matrix latents and a unitary predictor to imagine future states in partially observable environments. This model maintains joint-state invariants during blind rollout, learns action-sensitive dynamics through counterfactual targets, and demonstrates improved short-horizon information retention compared to LSTM-JEPA baselines.

Key insights

• Endowing a JEPA's latent space with an explicit belief structure (density matrix) and enforcing non-dissipative, unitary dynamics on its predictor enables robust imagination under partial observability.
• Learning action-conditioned future forecasting within JEPAs requires a shift from standard "teacher-forced" objectives to counterfactual targets that explicitly expose the predictor to diverse action outcomes.
• The improved performance of UWM-JEPA in blind rollout and action sensitivity is primarily due to its unique predictor dynamics and latent geometry, not merely a better initial observation encoder.

Summary by alphaXiv.