Search papers, labs, and topics across Lattice.
This paper critiques the limitations of traditional deterministic predictors in JEPA world models when faced with stochastic environments, where they fail to accurately represent the next latent state. The authors introduce MoP-JEPA, a method utilizing hard-assigned predictors that effectively quantize the transition distribution, allowing for improved planning performance. Experimental results demonstrate that MoP-JEPA significantly outperforms single-predictor rollouts and other state-of-the-art methods, achieving up to 0.85 success in planning tasks, while also employing a robust verification protocol to ensure accuracy.
Hard-assigned predictors in MoP-JEPA enable accurate planning in stochastic environments, achieving up to 42 times better performance than traditional methods.
JEPA world models predict the next latent state with a single deterministic predictor trained by latent regression. We show that this fails structurally when the environment is stochastic: at a branching transition, the regression-optimal predictor outputs the conditional mean of the successor embeddings, a point between the true next states that corresponds to no state at all. We prove 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 is the interface a planner consumes. On official OGBench offline data with leak-free evaluation, planning over single-predictor rollouts performs poorly ($0.02$--$0.09$ success) while planning over our predicted modes reaches up to $0.85$, ahead of deterministic, gated-MoE, and variational predictors on every task. Because multi-prediction evaluation invites coverage freeloading, a verification protocol is part of the method: an input-agnostic codebook control, a shuffled-context test, router-gated readouts, transition-precision guards, and a verified-route criterion in which the model proposes its transition graph blind and ground truth is used only to check the result. Under this criterion our method outperforms the strongest soft alternative on all three mazes ($2$--$5\times$), 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 the published OGBench baselines on the hardest maze. Multimodal dynamics decide whether a JEPA world model can plan at all; a mixture of predictors with hard assignment is a minimal and verifiable fix.
