Causal-JEPA: Learning World Models through Object-Level Latent Masking

World models require robust relational understanding to support prediction, reasoning, and control. While object-centric representations provide a useful abstraction, they are not sufficient to capture interaction-dependent dynamics. We therefore propose C-JEPA, a simple and flexible object-centric world model that extends masked joint embedding prediction from image patches to object-centric representations. By masking object-level latents and requiring each masked object state to be inferred from the surrounding context, C-JEPA imposes structured partial observability during training, creating counterfactual-like prediction queries that discourage shortcut solutions and make interaction-dependent prediction necessary under the learning objective. Empirically, C-JEPA leads to consistent gains in visual question answering, with an absolute improvement of about 20\% in counterfactual reasoning over the same architecture without object-level masking. On agent control tasks, C-JEPA enables substantially more efficient planning by using only 1\% of the total latent input features required by patch-based world models, while achieving comparable performance. Finally, we provide a formal analysis demonstrating that object-level masking induces useful inductive bias by controlling observability. Our code is available at https://github.com/galilai-group/cjepa.

To act intelligently in the world, AI systems need to predict how objects will move, interact, and affect one another. However, many current world models can make predictions without truly learning these object relationships. We propose C-JEPA, a model that learns by temporarily hiding information about some objects and predicting it from the surrounding objects. This encourages the model to focus on how objects relate to each other, similar to asking “what would this object be like if I could only see the rest of the scene?” In visual reasoning experiments, C-JEPA substantially improves counterfactual reasoning, gaining about 20 percentage points over the same architecture without object-level masking. In control tasks, it enables efficient planning while using far fewer input features than patch-based world models. This makes the model more compact while preserving strong performance. Overall, our results show that training world models to predict missing object information can improve their ability to reason about interactions, and our analysis explains why this masking strategy provides a useful learning bias.