Branch-JEPA is introduced, which replaces this point-valued transition with a context-weighted finite set of latent successors, and preserves more distinct futures, while full-set scoring improves the quality of the resulting predictive distribution.
Abstract
Joint-embedding predictive architectures (JEPAs) learn dynamics by predicting future observations in representation space. Yet most JEPA world models return one latent successor, even when hidden intent, partial observation, or stochastic dynamics make several futures plausible. We introduce Branch-JEPA, which replaces this point-valued transition with a context-weighted finite set of latent successors. Every branch is decoded independently, and the complete set is retained at inference. The architecture supports two complementary training regimes: specialization for recovering separated successors and full-set Energy-Score training for distributional fidelity. In a locked five-seed evaluation on the Argoverse~2 official validation split, full-set training improves trajectory Energy Score by $5.8$--$6.5\%$ and probability-weighted trajectory distance by $9.3$--$10.4\%$ over matched-$K{=}6$ assignment and transport objectives, while retaining $5.36$ endpoint-deduplicated effective branches. In a parameter-exact official-validation comparison, latent branching retains $10.3\%$ more effective modes and improves Energy Score, expected ADE, and Brier in all five paired seeds over branching only at the output decoder; every paired 95\% interval excludes zero. In an OGBench graph audit, Branch-JEPA increases teleport verified-route existence to $19.2\%$ versus $3.9\%$ for the MDN. Its raw-support advantage also persists with 29-D state and RGB observations. Together, latent branching preserves more distinct futures, while full-set scoring improves the quality of the resulting predictive distribution.
Latent world models plan by predicting how candidate actions advance learned latent dynamics. In self-predictive models, however, the encoder and predictor are optimized jointly and can co-adapt to latent transitions that are easy to predict but weakly constrained by the physical evolution of the scene. We introduce the cross-predictive JEPA (JEPA-x), which grounds latent dynamics in privileged physical trajectories. JEPA-x treats visual observations and physical states as corresponding views of the same action-conditioned trajectory, advances both through a shared predictor, and matches each prediction to the future representations of both modalities. This encourages the action-conditioned predictor to learn a common transition rule across the two views. Privileged physical state is used only during training, leaving a visual-only model at deployment. Empirical results show that JEPA-x reduces the rollout drift of a newly fitted predictor from $0.361$ to $0.104$ and increases mean control success from $53.6\%$ to $78.2\%$ on a multi-task suite spanning six evaluation subfamilies. We additionally show that direct physical-state regression improves decodability without improving forecastability or control, indicating that the benefit comes from shaping latent dynamics rather than merely encoding physical variables.
Kehan Wen, Ziming Li, Siyuan Luo et al.· 0 citations
Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting in representation space rather than reconstructing pixels, making them a natural backbone for latent model predictive control from offline demonstration logs. JEPA-style training optimizes short-horizon latent prediction, whereas planning requires a multi-step ranking of imagined futures by goal progress. Prior JEPA planners often inherit that ranking from embedding geometry, typically latent Euclidean distance, which arises as a byproduct of representation learning rather than as a progress cost mined from the logs. We propose Temporal-Distance-JEPA, which retains the LeWM encoder--predictor backbone and mines a directed temporal cost from reward-free trajectories: same-trajectory step order supplies positive targets, cross-trajectory pairs act as heuristic negatives, and a rollout-consistency term matches the planner horizon. The mined supervision serves two roles: as the deployed planning cost when progress is topological, and as a representation signal that improves Euclidean planning when contact geometry dominates. Under locked evaluation, deploying the mined cost raises Two-Room success to 100.0% versus LeWM's 97.4%, while shared Euclidean planning on the same temporally trained checkpoint raises OGB-Cube by 14.2 points over LeWM and improves Push-T. Against LeWM and the concurrent RC-aux baseline under locked evaluation, Temporal-Distance-JEPA matches or exceeds both methods on every environment. Ablations show that the directed head, cross-trajectory negatives, and rollout consistency each contribute. Temporal-Distance-JEPA narrows the train--plan gap for JEPA world-model planners by discovering temporal progress structure in offline logs and co-designing cost form with plan-time deployment. Code is available at https://github.com/HKBU-KnowComp/Temporal-Distance-JEPA.
Qantara is presented, an end-to-end JEPA whose joint training objective pairs a Brownian-bridge interpolant between consecutive clean latents on the state axis with noise-to-data flow matching on the action axis, and these results move JEPA world models from single-paradigm planners to multi-paradigm controllers.
Ruslan Rakhimov, George Bredis, Yuriy Maksyuta et al.· 2 citations
A latent world model trains its decoder on latents anchored to observations, then deploys it on the model's own free-running rollout, hundreds of steps past the last observation. Rollout-Decoded Reconstruction (RDR) closes this gap with a single loss term that free-runs the model during training exactly as evaluation will, decodes every rollout latent, and penalizes reconstruction error against ground truth. The term adds no parameters, costs training-time compute only, and reduces to the standard objective at weight zero, so every comparison in this paper is a one-flag A/B. On the chaotic Kuramoto-Sivashinsky equation, RDR raises valid prediction time (the time to first crossing of normalized error 0.5) from $3.87 \pm 0.23$ to $6.97 \pm 0.42$ time units at an identical 193,568 parameters, a $1.80\times$ improvement confirmed on seeds never used in selection and holding in 10 of 10 preregistered configurations at ratios of 1.71-2.50$\times$. The results come from a single system; a sweep in which the advantage grows with latent width is descriptive, and control experiments on two classic tasks are preliminary.
Method, a latent world-modeling framework based on orthogonal predictive factorization, is introduced, a latent world-modeling framework based on orthogonal predictive factorization that can be used by a readout, decoder, planner, or autoregressive rollout of an underlying system.
Contrastive inverse dynamics thus provides a distribution-free anti-collapse signal that requires no target network, stop-gradient, pretrained encoder, or reconstruction objective, and it is argued that the anti-collapse pressure can instead come from the transition data itself.