🤖 AI Summary
This work investigates how predictive representation learning objectives often discard exogenous features that, while unpredictable, are relevant for control due to their inherent bias toward predictability. Through a 2×2 controlled experimental design that independently manipulates feature controllability and control relevance, the study systematically evaluates six representation learning objectives—including JEPA, action-conditional JEPA, inverse dynamics models, and their reward- or controllability-augmented variants—on their ability to preserve such features. The analysis reveals, for the first time, the precise failure mechanism by which predictive objectives neglect control-relevant information. To address this, the authors propose Reward-Anchored JEPA, which leverages as little as 2% reward-labeled data. Empirical results demonstrate that all purely predictive, reward-free objectives fail to retain the feature (performing near chance), whereas the proposed method robustly recovers its representation across multiple environments and latent space dimensions.
📝 Abstract
Joint-embedding predictive (JEPA-style) objectives learn representations by predicting future latents. In doing so they can discard features that are exogenous (uncontrollable by the agent) yet control-relevant, even when those features are trivially encodable. This occurs because the objective optimizes temporal predictability rather than control-relevance. We isolate this failure mode in a controlled 2x2 experimental design that varies feature controllability and relevance independently, using a predictability knob that decouples a feature's temporal predictability from its control-relevance. Comparing six objectives: reconstruction, JEPA, action-conditioned JEPA, controllability-based JEPA, inverse dynamics under a random policy, and reward-grounded JEPA, we observe that all evaluated reward-free predictive objectives leave the exogenous control-relevant feature near chance accuracy, while a reward-grounded variant retains it selectively. The remedy is label-efficient and robust: as little as 2% of reward-labeled transitions recovers the feature, the effect holds across two environments with different surface forms, and it persists across latent dimensions from 16 to 1024. Comparing the learned latent geometry against bisimulation theory's prediction, the JEPA latent realizes only a small fraction of the class separation a supervised reference attains.