Standard JEPA world models predict the wrong future when environments can branch — and a new paper proves it.
Researchers published MoP-JEPA, an architecture that replaces a single deterministic predictor with a set of hard-assigned predictors, each covering one possible successor state. The structural problem it targets: when an environment reaches a branching transition, a conventional predictor outputs the conditional mean of the successor embeddings — a point between the real next states that corresponds to no real state at all. The paper proves this collapse applies to gated mixture-of-experts designs too, not just naive single-predictor setups. Hard assignment fixes it by routing each predictor to a distinct mode of the transition distribution, turning the model's output into a map of real futures a planner can actually use.
The practical gap is hard to ignore. On OGBench offline benchmark tasks, planning over single-predictor rollouts scored between 0.02 and 0.09 success rate. Planning over MoP-JEPA's predicted modes reached 0.85, outperforming deterministic, gated-MoE, and variational baselines on every task. Because multi-prediction setups can game metrics by claiming coverage they did not earn, the paper includes a verification protocol — codebook controls, shuffled-context tests, and a verified-route criterion where the model proposes its transition graph without seeing ground truth. Under that stricter standard, MoP-JEPA outperformed the best soft alternative by two to five times across three maze tasks.
The same model, run live in the real environment, placed second among seven baselines on OGBench's hardest maze — a useful proof of concept, if not a rout.