Predictive Objectives Discard Exogenous Control-Relevant Features: A Controlled Mechanistic Study

Technical Summary: Predictive Objectives Discard Exogenous Control-Relevant Features

Problem Statement

This study investigates a specific failure mode in Joint-embedding Predictive Architectures (JEPA) and similar self-prediction objectives. While these models learn representations by predicting future latent states, they optimize for temporal predictability rather than control-relevance. The paper posits that this optimization bias causes the model to discard features that are exogenous (uncontrollable by the agent) yet control-relevant (determinative of the optimal action or reward), even when those features are trivially encodable.

The core tension lies between two theoretical frameworks:

1. Self-Prediction: Discards "unpredictable detail" to achieve scalability and robustness against nuisance textures.
2. Bisimulation Theory: Argues that optimal control representations must preserve features that distinguish states by their reward-relevant futures, regardless of their predictability.

The paper focuses on a specific quadrant of a feature taxonomy: features that are uncontrollable (exogenous) but relevant. Prior work has extensively studied controllable features and irrelevant exogenous distractors, but the "uncontrollable-relevant" case has not been isolated and measured as a distinct failure mode.

Methodology

The authors employ a controlled mechanistic study using a 2 × 2 experimental design to independently vary feature controllability and relevance.

• Environments: The primary experiments use synthetic environments (QuadrantEnv, SwitchColorEnv, GridWorldHiddenRuleEnv) with varying surface forms (resolutions, textures) but identical underlying logic.
  • Feature Dynamics: Features are defined by two axes:
    • Controllability: Whether the agent's action influences the feature's next state.
    • Relevance: Whether the feature determines the reward (and thus the optimal action).
  • The Predictability Knob: A parameter ($p_{repeat}$) controls the temporal predictability of exogenous features. At $p_{repeat} = 0.5$, the feature is an i.i.d. Bernoulli draw (maximally unpredictable); at $p_{repeat} = 1.0$, it is constant. Crucially, control-relevance is held fixed across the entire range, decoupling predictability from relevance.
• Objectives Compared: The study evaluates six objectives plus a supervised reference:
  1. Reconstruction: Pixel reconstruction (baseline).
  2. JEPA: Latent self-prediction (baseline).
  3. Action-Conditioned JEPA: Predicts future latent conditioned on action.
  4. Controllability-Based JEPA: Uses an inverse-dynamics head to predict actions.
  5. Inverse Dynamics (Random Policy): Predicts actions under a random policy.
  6. Reward-Grounded JEPA: A proposed variant incorporating reward labels.
  7. Supervised Reference: Trained on ground-truth feature labels (upper bound).
• Metrics:
  • Linear Probe Accuracy: Measures if the feature can be decoded from the frozen latent.
  • InfoNCE Mutual Information: Estimates $I(Z; c)$ to confirm information retention without shared decision boundaries.
  • Effective Rank: Ensures the representation has not collapsed (low rank) but rather selectively dropped the feature.

Key Results

The experiments reveal a stark divergence between reward-free predictive objectives and reward-grounded methods regarding the exogenous control-relevant feature (Cell 4).

1. Failure of Reward-Free Predictive Objectives:

   • All evaluated reward-free objectives (JEPA, Action-Conditioned JEPA, Controllability-Based JEPA, and Inverse Dynamics under random policy) fail to retain the exogenous control-relevant feature.
   • Performance: Linear probe accuracy hovers near chance ($\approx 0.51$), and Mutual Information is near zero ($\approx 0$ nats).
   • Selectivity: This is not a representational collapse. The effective rank remains high (38–42 in QuadrantEnv), indicating the representation is rich but has selectively discarded this specific feature.
   • Capacity Independence: Increasing latent dimensions from 16 to 1024 does not rescue the feature for JEPA; the failure persists, confirming it is an objective-level structural issue, not an architectural capacity bottleneck.

2. Success of Reward Grounding:

   • The Reward-Grounded JEPA variant selectively retains the exogenous control-relevant feature, achieving linear probe accuracy of $1.00$ and Mutual Information of $\approx 0.693$ nats (the theoretical maximum for a one-bit feature).
   • Label Efficiency: The remedy is highly sample-efficient. As little as 2% of transitions labeled with rewards is sufficient to recover the feature (accuracy $\approx 0.77$).
   • Selectivity: Unlike reconstruction (which retains all features, including irrelevant distractors), the reward-grounded variant retains only the relevant cells (Controllable-Relevant and Exogenous-Relevant) and discards irrelevant ones.

3. Bisimulation Comparison:

   • Theoretical bisimulation distance for the exogenous control-relevant feature is $1.0$ (states should be distinct).
   • The supervised reference and reconstruction baselines realize a latent class separation of $\approx 1.998$.
   • JEPA realizes only $\approx 0.105$ separation, confirming that the latent geometry fails to encode the reward-relevant distinction predicted by bisimulation theory.

Significance and Claims

The paper explicitly states its aim is not to benchmark world models on natural tasks or to claim that all predictive objectives are flawed for general representation learning. Instead, the significance lies in:

• Isolation of a Failure Mode: Providing the first controlled isolation and measurement of the "exogenous control-relevant" failure mode, demonstrating that predictability and control-relevance can be orthogonal.
• Mechanistic Insight: Showing that purely predictive objectives optimize for temporal predictability, which can lead to the systematic loss of information critical for control if that information is unpredictable.
• Selective Remedy: Demonstrating that a small amount of reward signal (2% of transitions) can selectively recover this lost information without requiring full supervision or sacrificing the robustness of predictive objectives to irrelevant noise.
• Theoretical Gap: Quantifying the gap between the class separation required by bisimulation theory and the separation actually realized by JEPA latents.

The authors conclude that while predictive objectives are effective for their designed goals (ignoring nuisance texture), they structurally risk discarding control-relevant information that is temporally unpredictable, a risk that can be mitigated by minimal reward grounding.