Latent world models predict action effects by intervening on reality
Do-JEPA: From Masking to Intervention in Latent World Models
Computer Vision and Pattern RecognitionArtificial IntelligenceMachine Learning
Summary
It is hard for computer models to know what part of the future changes truly come from their actions rather than just happening alongside them. The authors designed a new way to train these models by comparing what actually happens in a simulation when they take an action versus when they do not, focusing on the difference caused by the action itself. This helps the model understand cause and effect better by looking at changes in the world, not just what it sees. Their tests show this method improves prediction accuracy of real action effects from pixels and simulated physics environments.
What this means in practice
- •For robotics engineers: Improve robot prediction accuracy of how their actions physically change the environment using models trained on simulated interventions.
- •For game ai developers: Build more accurate AI agents that understand what their moves cause in the game world by comparing simulated outcomes under different actions.
Authors
Hossein Resani, Javen Qinfeng Shi
Abstract
Latent world models are trained to predict what happens next, so nothing in their objective separates what an action caused from what merely co-occurred with it. Object-masking models such as C-JEPA intervene on what the predictor can see; we intervene on what physically happens. From one saved simulator state we run the dynamics under an action $a$ and under a reference action $a_{\varnothing}$, and train the model to predict the difference $Δz=z^{a}-z^{a_{\varnothing}}$ between the two latent futures. The resulting objective, Do-JEPA, has an effect loss, a support loss (where the action enters), a propagation loss (where its effect travels) and invariance losses (what must not change). In a synthetic system with object-aligned variables, support supervision finds the directly intervened object in 99.95% of test cases, where a sparse action mask sends the action to a nuisance slot in every case, and response-onset supervision recovers the ring-shaped propagation graph (edge AUROC 0.975 vs. 0.624). From pixels, the effect loss beats a control trained on exactly the same data: it lowers latent effect error by 28.4% on an end-to-end LeWM model and physical effect error by 13.5% when trained and tested on natural action sequences, and on three independently generated CausalWorld benchmarks it lowers responsive effect error by about 20% under physics shifts and the latent context sensitivity of predicted effects by 66%. Trained from scratch it costs factual accuracy; fine-tuning an existing model with it removes this cost. Together, these results show that intervening on the world, rather than on what the model sees, helps latent world models predict what their actions cause.