From TD Learning to JEPA: Why Future-Latent Prediction Is the Real Lineage

The thesis

A lot of modern "world model" research can be reduced to one surprisingly old idea:

A good representation of the present is one that makes the future predictable.

That sentence sounds vague until we formalize it. In reinforcement learning, the agent sees observations $o_t$ , takes actions $a_t$ , receives rewards $r_t$ , and transitions through hidden states $s_t$ . The core problem is not just learning a policy. The core problem is learning a state representation:

z_t = f_\theta(o_{\leq t}, a_{<t})

such that $z_t$ preserves the information needed for future prediction and control.

The modern JEPA family did not appear out of nowhere. It inherits from a long lineage:

TD learning → successor representations → predictive state representations
→ contrastive predictive coding → BYOL / SimSiam / DINO
→ PBL / SPR in reinforcement learning
→ I-JEPA / V-JEPA / action-conditioned JEPA world models

The key shift across this lineage is:

predict pixels  →  predict rewards/values  →  predict future observations
→ predict future latent representations

The deep idea is not "reconstruction." The deep idea is bootstrapped prediction in representation space.

1. TD learning: prediction by bootstrapping from your own future prediction

Temporal-difference learning is the original bootstrapping idea in RL. Instead of waiting until the full return is observed, TD updates a prediction using the next prediction. Sutton's 1988 paper introduced TD methods as prediction-learning procedures that assign credit using the difference between temporally successive predictions.

For a value function $V_\theta(s)$ , the TD target is:

y_t = r_t + \gamma V_{\bar{\theta}}(s_{t+1})

and the loss is:

\mathcal{L}_{TD}(\theta) = \left( V_\theta(s_t) - \left[ r_t + \gamma V_{\bar{\theta}}(s_{t+1}) \right] \right)^2

The philosophical move is huge:

You can learn from a target produced by your own model, as long as the target is more stable than the thing being updated.

That is the same shape that later appears in BYOL, SPR, and JEPA.

TD:
  current value predicts reward + future value

BYOL:
  current representation predicts target-network representation

SPR:
  current latent predicts future target-network latent

JEPA:
  context latent predicts missing/future target latent

TD bootstraps values. BYOL/SPR/JEPA bootstrap representations.

2. Successor representations: states should be similar if their futures are similar

The successor representation made the next conceptual jump: the meaning of a state is not just its current pixels or coordinates, but the distribution of future states it leads to. Dayan's 1993 successor representation paper argues that appropriate generalization between states is determined by how similar their successors are.

For a policy $\pi$ , define a feature vector $\phi(s)$ . The successor representation is:

M^\pi(s) = \mathbb{E}_\pi \left[ \sum_{k=0}^{\infty} \gamma^k \phi(s_{t+k}) \mid s_t = s \right]

If rewards are linear in features:

r(s) = \phi(s)^\top w

then value becomes:

V^\pi(s) = M^\pi(s)^\top w

This is the first-principles bridge to representation learning:

Two observations should be close in representation space if they imply similar future trajectories.

This is already JEPA-like, except it is written in classical RL notation instead of deep latent prediction notation.

3. Predictive state representations: state is a vector of future predictions

Predictive State Representations, or PSRs, made the idea even sharper: instead of representing hidden state directly, represent the system through predictions about future observable events. Littman, Sutton, and Singh showed that dynamical system states can be represented by multi-step, action-conditional predictions of future observations.

A PSR state can be written as:

q_t = \left[ P(\tau_1^O \mid h_t, \tau_1^A), \; P(\tau_2^O \mid h_t, \tau_2^A), \; \ldots, \; P(\tau_n^O \mid h_t, \tau_n^A) \right]

where $h_t$ is the history, $\tau_i^A$ is a future action sequence, and $\tau_i^O$ is a future observation sequence.

In plain English:

The state is not a hidden object behind the observations. The state is the set of future predictions you can make from the current history.

That is almost the entire philosophical foundation of world models.

4. CPC: predict the future in latent space, not pixel space

Contrastive Predictive Coding moved predictive representation learning into deep learning. CPC learns representations by predicting future latent samples with an autoregressive model and a contrastive loss. The authors describe the key idea as learning representations by predicting the future in latent space using powerful autoregressive models.

Let:

z_t = f_\theta(x_t)

c_t = g_\psi(z_{\leq t})

CPC tries to identify the true future latent $z_{t+k}$ among negatives $\tilde{z}$ :

\mathcal{L}_{CPC} = -\log \frac{ \exp(z_{t+k}^\top W_k c_t) }{ \sum_{\tilde{z} \in \mathcal{N}} \exp(\tilde{z}^\top W_k c_t) }

This is where the future-latent objective becomes explicit:

Do not reconstruct x_{t+k}.
Encode x_{t+k} into z_{t+k}.
Predict z_{t+k}.

But CPC still uses negatives. Later methods ask: can we bootstrap without negatives?

5. BYOL: bootstrapping representations without negatives

BYOL introduced the clean modern bootstrap representation recipe. It uses an online network and a target network. The online network predicts the target-network representation of another augmented view of the same image; the target network is updated by an exponential moving average of the online network.

Let $x_1 = t_1(x)$ and $x_2 = t_2(x)$ be two augmentations.

Online branch:

z_1 = g_\theta(f_\theta(x_1))

p_1 = q_\theta(z_1)

Target branch:

z_2^- = g_{\bar{\theta}}(f_{\bar{\theta}}(x_2))

Loss:

\mathcal{L}_{BYOL} = \left\| \text{norm}(p_1) - \text{sg}(\text{norm}(z_2^-)) \right\|_2^2

Target update:

\bar{\theta} \leftarrow \tau \bar{\theta} + (1-\tau)\theta

The important tricks are:

online network gets gradient
target network gets no gradient
target network moves slowly
predictor exists on online side
representations are normalized

This solves the "where do labels come from?" problem by making the model's own slow-moving representation into the label.

The scary failure mode is collapse:

f_\theta(x) = c \quad \forall x

BYOL-style methods avoid this because the training dynamics, asymmetry, predictor, stop-gradient, and target update prevent the trivial constant solution in practice. SimSiam later showed that even without negative pairs, large batches, or momentum encoders, a stop-gradient operation plays an essential role in preventing collapse.

This is the exact mechanism later reused in RL latent prediction and JEPA-style systems.

6. PBL: bootstrapped latent prediction enters deep RL

Predictions of Bootstrapped Latents, or PBL, brought this idea into multitask RL. PBL learns representations by predicting latent embeddings of future observations, focusing on structured information about environment dynamics.

A simplified PBL objective:

z_t = f_\theta(o_t)

z^-_{t+k} = f_{\bar{\theta}}(o_{t+k})

\hat{z}_{t+k} = m_\psi(z_t, a_t, \ldots, a_{t+k-1})

\mathcal{L}_{PBL} = \sum_{k=1}^{K} \left\| \hat{z}_{t+k} - \text{sg}(z^-_{t+k}) \right\|^2

This is the important conceptual step:

The agent is not just learning a visual representation. It is learning a representation whose future evolution is predictable.

That is much closer to control than static image SSL.

7. SPR: the clean RL version of future-latent bootstrapping

Self-Predictive Representations, or SPR, is the paper I would point to if someone asked, "what is the direct ancestor of action-conditioned JEPA for RL?" SPR trains an agent to predict its own latent state representations multiple steps into the future, using an EMA target encoder and a learned transition model. It reported a 55% relative improvement over prior Atari 100k state of the art in median human-normalized score.

SPR can be written as:

z_t = f_\theta(o_t)

\hat{z}_{t+1} = m_\psi(z_t, a_t)

\hat{z}_{t+2} = m_\psi(\hat{z}_{t+1}, a_{t+1})

\vdots

z^-_{t+k} = f_{\bar{\theta}}(o_{t+k})

\mathcal{L}_{SPR} = \sum_{k=1}^{K} \left\| \text{norm}(\hat{z}_{t+k}) - \text{sg}(\text{norm}(z^-_{t+k})) \right\|_2^2

This objective is trained alongside the RL loss:

\mathcal{L} = \mathcal{L}_{RL} + \lambda \mathcal{L}_{SPR}

This is basically BYOL, but temporally extended and made action-aware.

BYOL:
  x under augmentation 1 → predict x under augmentation 2 in latent space

SPR:
  o_t plus actions → predict o_{t+k} in latent space

This is why I think modern JEPA researchers are not "stealing" from SPR in a dirty way, but they are absolutely inheriting the same core principle:

predict future or missing latent representations, not pixels.

8. Theoretical evidence: why self-predictive learning does not have to collapse

The obvious objection to BYOL/SPR-style learning is:

f(o) = c

If every observation maps to the same constant vector, future-latent prediction becomes trivial. So why does it work?

The 2022 paper "Understanding Self-Predictive Learning for Reinforcement Learning" studies this directly. It identifies that faster predictor optimization and semi-gradient updates are key for avoiding collapse, and shows in an idealized setup that self-predictive learning performs spectral decomposition on the state transition matrix, capturing transition-dynamics information.

That matters because it turns the vague claim into a real mechanism:

self-predictive losses can discover the slow/dominant structure of environment dynamics.

A later 2024 paper on action-conditional self-predictive RL analyzes BYOL-AC and connects these objectives to low-rank approximations of dynamics and to value, Q-value, and advantage-like functions. It also reports that the action-conditioned version performs better overall across several empirical settings.

So the theory timeline looks like this:

1988 TD:
  bootstrapping predictions can learn future outcomes

1993 successor representations:
  good state similarity is future-occupancy similarity

2001 PSRs:
  state can be represented as action-conditional predictions

2020 BYOL:
  self-generated target representations can train visual encoders without negatives

2020 PBL / SPR:
  future-latent prediction improves sample-efficient visual RL

2022-2024 self-predictive theory:
  these objectives can recover transition structure and connect to value/Q/advantage quantities

2023-2025 JEPA:
  scale latent prediction to images, video, and action-conditioned world models

9. JEPA: the same principle becomes a general architecture

JEPA stands for Joint-Embedding Predictive Architecture. LeCun's 2022 position paper framed JEPA as part of a larger architecture for autonomous machine intelligence, combining predictive world models, intrinsic motivation, and hierarchical joint-embedding architectures trained with self-supervised learning.

The generic JEPA formulation is:

Context encoder:

z_c = f_\theta(x_c)

Target encoder:

z_y^- = f_{\bar{\theta}}(x_y)

Predictor:

\hat{z}_y = p_\psi(z_c, m)

Loss:

\mathcal{L}_{JEPA} = D \left( \hat{z}_y, \text{sg}(z_y^-) \right)

where $x_c$ is visible context, $x_y$ is the masked or future target, $m$ describes the mask or target position, and $D$ is usually a latent-space distance.

The important part is what JEPA does not do:

\hat{x}_y \neq \text{decoder}(z_c)

It does not need to reconstruct pixels. It predicts embeddings.

I-JEPA applies this to images: from a context block, predict representations of target blocks in the same image. The I-JEPA paper emphasizes that it is non-generative and learns semantic image representations without relying on hand-crafted augmentations.

V-JEPA applies this to video: learn visual representations by feature prediction from video, without pretrained image encoders, text, negative examples, reconstruction, or other supervision.

V-JEPA 2 scales the idea further: it pretrains on over 1 million hours of internet video and then post-trains an action-conditioned latent world model, V-JEPA 2-AC, using less than 62 hours of robot video, enabling zero-shot image-goal planning on Franka arms in two labs.

So the lineage is not accidental:

SPR:
  z_t, a_t, ..., a_{t+k-1} → z^-_{t+k}

V-JEPA:
  video context → future/masked video features

V-JEPA 2-AC:
  latent state + action → future latent state for planning

For agents, the most important version is not passive I-JEPA. It is action-conditioned JEPA:

\hat{z}_{t+1} = p_\psi(z_t, a_t)

\mathcal{L}_{AC\text{-}JEPA} = \left\| \hat{z}_{t+1} - \text{sg}(f_{\bar{\theta}}(o_{t+1})) \right\|^2

That is basically the SPR/PBL control form, now scaled and reframed as world modeling.

10. Now the encoder-sharing question

The practical question is:

In a vision-action RL model, should the actor and critic share the encoder?

The options are usually:

A. actor and critic have separate visual encoders

B. actor and critic share one encoder,
   but the shared encoder is updated only by policy loss

C. actor and critic share one encoder,
   but the shared encoder is updated only by critic/value loss

D. actor and critic share one encoder,
   the shared encoder is updated by both policy and critic losses

E. shared low-level encoder,
   separate actor/critic trunks,
   shared encoder updated by critic/value + future-latent predictive loss

My pick is E.

Not because it is fashionable. Because it falls out of the objective geometry.

11. First-principles derivation

An actor-critic agent has at least two different problems.

The actor wants:

\pi_\omega(a_t \mid z_t)

and is updated by something like:

\nabla_\omega J = \mathbb{E} \left[ \nabla_\omega \log \pi_\omega(a_t \mid z_t) \, A^\pi(s_t, a_t) \right]

The critic wants:

V_\nu(z_t) \quad \text{or} \quad Q_\nu(z_t, a_t)

and is updated by:

\mathcal{L}_{critic} = \left( Q_\nu(z_t, a_t) - \left[ r_t + \gamma Q_{\bar{\nu}}(z_{t+1}, a_{t+1}) \right] \right)^2

These two objectives are related, but they do not ask for the same representation.

The actor wants features that discriminate which action to take.

The critic wants features that predict what will happen and how much return follows.

A recent ICLR 2025 study found that when actor and critic representations are separated, they specialize: actor representations focus more on action-relevant information, while critic representations specialize in value and dynamics information.

So fully sharing everything is theoretically suspicious.

But fully separating everything is also expensive:

separate actor encoder:
  must learn vision for action selection

separate critic encoder:
  must learn vision for value/dynamics

result:
  doubled visual learning burden

For pixels, that is painful. Vision is the expensive part.

So the first-principles compromise is:

share early perception
separate late semantics
do not force actor and critic to use the exact same high-level representation

Architecture:

pixels o_t
  ↓
shared visual stem fθ
  ↓
h_t
  ├── actor trunk uω(h_t) → π(a_t | h_t)
  ├── critic trunk vν(h_t) → Q(h_t, a_t)
  └── predictive head mψ(h_t, a_t) → ĥ_{t+1}

Loss:

\mathcal{L} = \mathcal{L}_{actor} + \mathcal{L}_{critic} + \lambda \mathcal{L}_{pred}

But gradient routing matters:

\nabla_\theta \approx \nabla_\theta \mathcal{L}_{critic} + \lambda \nabla_\theta \mathcal{L}_{pred}

and usually:

\nabla_\theta \mathcal{L}_{actor} \approx 0

at least at the start.

In code terms:

h = encoder(obs)

# critic update
q = critic(h, action)
critic_loss.backward()      # updates critic + encoder

# predictive update
pred = transition(h, action)
pred_loss.backward()        # updates predictive head + encoder

# actor update
h_detached = h.detach()
policy_loss = actor_loss(actor(h_detached), critic)
policy_loss.backward()      # updates actor, not encoder

DrQ-v2, a strong visual continuous-control baseline, follows this general spirit: the encoder is optimized during critic update, and the actor update uses detached encoded observations in the official implementation. DrQ-v2 also showed strong results on DeepMind Control Suite, including pixel-based humanoid control, while remaining simple and computationally efficient.

This gives us a clean rule:

Let the critic and predictive objective teach the shared visual stem what the world is. Let the actor learn what to do with that representation.

12. Why policy-only encoder updates are weak

Policy gradients are high variance:

\nabla_\theta \log \pi(a_t \mid f_\theta(o_t)) \, A_t

This signal does not directly say:

where is the object?
what moved?
what is controllable?
what will happen after action a?

It only says:

make this sampled action more or less likely under this noisy advantage estimate

That is a bad primary teacher for vision.

So this option is usually weak:

shared encoder updated only by policy loss

It asks sparse, noisy action-selection pressure to solve perception.

13. Why critic-only is a strong baseline but not the whole answer

Critic loss is denser:

r_t + \gamma Q(z_{t+1}, a_{t+1})

The critic repeatedly sees consequences. It is closer to world modeling than the actor is. So critic-only encoder updates are often a strong baseline.

But critic-only can become value-myopic. If a visual feature is not immediately useful for value prediction under the current policy, the critic may ignore it, even if it is useful for future exploration or action selection.

That is why the future-latent predictive loss matters.

\mathcal{L}_{pred} = \sum_{k=1}^{K} \left\| \text{norm}(\hat{z}_{t+k}) - \text{sg}(\text{norm}(z^-_{t+k})) \right\|^2

This forces the encoder to preserve dynamics-relevant information even before the reward function fully reveals why that information matters.

14. The result I would bet on

The best general configuration for vision-action models is:

shared low-level encoder
separate actor and critic trunks
encoder updated by critic/value loss
encoder updated by future-latent predictive loss
actor gradient detached from shared encoder initially
optional small actor-gradient flow later

Formally:

h_t = f_\theta(o_t)

\pi_t = \pi_\omega(u_\omega(h_t))

Q_t = Q_\nu(v_\nu(h_t), a_t)

\hat{h}_{t+k} = m_\psi(h_t, a_t, \ldots, a_{t+k-1})

h^-_{t+k} = f_{\bar{\theta}}(o_{t+k})

\mathcal{L}_{total} = \mathcal{L}_{actor} + \mathcal{L}_{critic} + \lambda \sum_{k=1}^{K} \left\| \text{norm}(\hat{h}_{t+k}) - \text{sg}(\text{norm}(h^-_{t+k})) \right\|^2

with gradient routing:

\nabla_\theta \mathcal{L}_{total} = \nabla_\theta \mathcal{L}_{critic} + \lambda \nabla_\theta \mathcal{L}_{pred} + \alpha \nabla_\theta \mathcal{L}_{actor}

where I would start with:

\alpha = 0

and maybe later test:

\alpha \in [0.01, 0.1]

This is the first-principles reason behind the pick:

policy loss:
  best for action selection
  bad primary vision teacher

critic loss:
  better dense signal
  learns value/dynamics-relevant features

future-latent prediction:
  learns dynamics structure even before reward is informative

shared early encoder:
  sample-efficient perception

separate trunks:
  actor and critic can specialize

15. The actual research hypothesis

The blog-worthy hypothesis is not merely:

shared encoder or not?

The better hypothesis is:

In vision-action agents, the optimal representation is not fully shared or fully separate. Early perception should be shared because pixels are expensive. But actor, critic, and world-model objectives should separate after the visual stem because they impose different geometries on latent space.

Or more bluntly:

Share pixels. Do not blindly share objectives.

This follows directly from the lineage.

TD taught us to bootstrap predictions.

Successor representations taught us that state similarity should be future-similarity.

PSRs taught us that state can be a vector of future predictions.

CPC taught us to predict futures in latent space.

BYOL taught us that the targets can come from a slow-moving copy of ourselves.

SPR taught us that this is powerful for RL from pixels.

JEPA scaled the principle into a general architecture for world modeling.

So for a modern vision-action model, the clean architecture is not:

one encoder, all gradients smashed together

It is:

one shared perceptual stem
separate actor/critic/world-model pathways
future-latent prediction as the representation teacher
careful gradient routing

That is the configuration I would start with, and the one I would expect to generalize beyond Atari or robot benchmarks into browser agents, self-driving, VLAs, and any system where vision must support both action and prediction.

The deeper conclusion is simple:

Intelligence is not built by reconstructing the present. It is built by compressing the present into the variables that make the future controllable.