Learning latent dynamics is essential for building action‑conditioned world models that support planning. Among existing approaches, Joint Embedding Predictive Architectures (JEPA) provide a natural framework by modeling state transitions directly in latent space rather than at the pixel level. Nevertheless, JEPA‑based world models suffer from two drawbacks: (i) physical grounding does not explicitly aggregate temporal context, and (ii) action planning relies heavily on generic sampling distributions.
To address these issues we introduce $\mathcal{G}$-JEPA. During training, a shared causal temporal module aggregates temporal context across observation and prediction sequences under supervision of physical states, grounding the latent representations in task‑relevant physical variables.
During planning, $\mathcal{G}$-JEPA performs action inference by stepwise conditioning a generative action distribution on the current latent, the goal latent, and their evolving discrepancy. This yields anchored action proposals that guide model‑based online refinement, making the search through the learned world model more directed.
Empirical results on manipulation benchmarks demonstrate improved planning performance, especially as the temporal gap between initial and goal observations grows. The code is publicly available at https://github.com/myl11045/G-JEPA.
Review