Counterfactual JEPAs for World Models in the Real World
Randall Balestriero, Assistant Professor of Computer Science, Brown University; researcher at AMI Labs
Abstract: This talk presents a sequence of recent results on Joint-Embedding Predictive Architectures (JEPAs), spanning self-supervised representation learning, action-conditioned world models, and counterfactual reasoning. We begin with LeJEPA, which addresses representation collapse through a distributional principle: embeddings are regularized toward an isotropic Gaussian using Sketched Isotropic Gaussian Regularization, producing a simple predictive objective with theoretical guarantees and without teacher–student networks or stop-gradient heuristics. We then introduce LeWorldModel, which extends this construction to sequential data by learning an encoder and action-conditioned latent transition model, (z_{t+1} \approx f_\theta(z_t,a_t)), end-to-end from raw observations. By combining next-embedding prediction with Gaussian latent regularization, the model supports planning and control directly in representation space, without reconstructing future pixels. Finally, we discuss new directions in counterfactual JEPAs, where object-level masking and latent interventions create structured partial-observability problems that force the predictor to model interaction-dependent dynamics rather than exploit local correlations. Throughout, we emphasize the geometric and statistical principles connecting these methods, and ask when latent prediction is sufficient for learning causal, controllable abstractions of a dynamical world.
Cost: free
Kisa Kowal
(847) 491-3974
Email