PlanJEPA: Learning Control-Structured Representations for Latent Planning

Learning representations that support planning is essential for latent world models. However, existing representation-learning objectives focus on latent-state prediction or action recoverability and do not constrain the complexity of the control relation between actions and latent transitions, so an accurate predictor can still operate over a latent space with highly nonlinear control relations. Consequently, accurate latent prediction does not necessarily lead to successful planning, since goal costs, candidate-trajectory ranking, and gradient-based action optimization all depend on the control structure exposed by the latent space. We introduce PlanJEPA, which learns control-structured representations with a multi-horizon linear action objective. By requiring actions to be predicted through a shared linear map of current and future latent states across multiple temporal horizons, PlanJEPA encourages the encoder to organize control-relevant variations directly in the representation. We prove that the imposed linear action relation forces nontrivial control errors to induce nontrivial latent goal errors and improves the local conditioning of goal-directed planning objectives. Empirically, PlanJEPA better organizes control-relevant state changes in latent space and consistently outperforms strong latent world model baselines across closed-form control, learned policies, cross-entropy method search, and gradient-based planning.