Abstract
JEPA world models predict the next latent state with a single deterministic predictor trained by latent regression. This approach structurally fails in stochastic environments: at branching transitions, the regression-optimal predictor outputs the conditional mean of the successor embeddings, a point that lies between the true next states, effectively corresponding to a non-existent state. We demonstrate this collapse for deterministic and gated mixture-of-experts predictors and prove that MoP-JEPA's hard-assigned predictors converge instead to a quantizer of the transition distribution: one head per successor mode, enumerable in a single forward pass, which serves as the interface consumed by a planner.
On official OGBench offline data, planning over single-predictor rollouts performs poorly ($0.02$--$0.09$ success) while planning over our predicted modes achieves up to $0.85$, surpassing deterministic, gated-MoE, and variational predictors on every task.
To address coverage freeloading in multi-prediction evaluation, our method includes a verification protocol: input-agnostic codebook control, shuffled-context tests, router-gated readouts, transition-precision guards, and a verified-route criterion where the model proposes its transition graph blind and ground truth is used solely to check results. Under this criterion, our method outperforms the strongest soft alternative on all three mazes ($2$--$5 imes$), and the protocol identifies the remaining gap in that baseline's raw scores as routes through predicted transitions that do not exist.
The same model executes in the real environment, placing second of seven against published OGBench baselines on the hardest maze. Multimodal dynamics determine whether a JEPA world model can plan at all; a mixture of predictors with hard assignment is a minimal and verifiable fix.
Blogger's Review: This paper showcases the potential of MoP-JEPA in stochastic environments by enhancing planning capabilities through a mixture of hard-assigned predictors. This advancement not only improves success rates but also provides a validated approach for future model designs, making it worth attention!