REVIEW 3 major objections 5 minor 10 references
JEPA for RL: Investigating Joint-Embedding Predictive Architectures for Reinforcement Learning
T0 review · 3 major / 5 minor · reviewed 2026-08-16 · deepseek-v4-flash
Pith's one-line read This paper shows that Joint-Embedding Predictive Architectures can be adapted to reinforcement learning, with a variance regularizer allowing them to learn state representations without reward-gradient feedback.
desk verdict A sensible, clearly-written JEPA-for-RL adaptation with an instructive collapse ablation, but the central claim is under-supported by one narrow environment and a missing control. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The machinery is the JEPA pipeline: a vision-transformer x-encoder that takes the three frames $f_{t-2}, f_{t-1}, f_t$ with positional encodings for both patch location and frame index; a target y-encoder with weights updated as a moving average of the x-encoder, applied to frames $f_{t-1}, f_t, f_{t+1}$; a two-layer predictor that is fed the action to output a predicted embedding; and the squared Euclidean loss between predicted and target embeddings. Collapse prevention is the second piece of machinery: when the RL losses cannot reach the encoder, a variance regularizer encourages the batch-wise embedding variance to stay above zero, and the paper reports this restores learning.
What would settle it
Run the paper's regularized no-gradient configuration on an Atari game such as Pong with the same architecture and check whether the batch-wise embedding variance stays above $10^{-7}$ and the episodic return improves; if the agent fails or collapses despite the variance regularizer, the central claim is refuted.
Extended reading notes
Core claim
On the paper's own terms, the discovery is that a JEPA trained purely by latent prediction can carry representation learning for RL from images. The evidence is a four-way ablation on Cart Pole: the baseline encoder trained only by actor-critic gradients learns limited behavior; adding the JEPA loss improves reward and avoids plateau; stopping RL gradient flow while keeping the JEPA loss leads to model collapse, with batch-wise embedding variance dropping below $10^{-7}$; and adding a variance regularizer prevents that collapse so the actor can learn from the embeddings. The paper explicitly concludes that JEPA is able to learn informative state representations without gradient propagation from the reinforcement learning task, and that JEPAs are promising candidates for reinforcement learning.
Load-bearing premise
The load-bearing premise is that Cart Pole, tested with five runs and no error bars, stands in for all low-dimensional-state image-based RL tasks, so if Cart Pole is not representative of harder Atari-style environments, the paper's conclusion that JEPAs are promising for RL does not follow.
Editorial extensions
If this is right
- If JEPA plus a variance regularizer can replace reward-driven encoder training, pixel-based RL agents could learn reusable state representations before seeing any reward, cutting training cost.
- Combining the JEPA loss with actor-critic gradient flow accelerates early reward growth and avoids the plateau seen in the baseline on Cart Pole.
- Because the representation is a single 64-dimensional classification-token embedding, the approach is matched to low-dimensional-state tasks such as Atari games, as the authors intend.
- Batch-wise embedding variance is a cheap, observable collapse indicator that can be monitored during RL training to detect representation failure early.
Reading between the lines
- A natural extension not tested in the paper: running the same four configurations on Atari games with strong visual distractors would test whether the regularized no-gradient variant still avoids collapse and learns, or whether harder perception demands RL-gradient feedback.
- If JEPA representations are learned purely from next-frame latent prediction, they may transfer across reward functions defined on the same environment, a property the paper does not investigate.
- Using a single classification token rather than all patch embeddings keeps the representation low-dimensional but may discard spatial details needed for fine-grained control; the paper does not measure this trade-off.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes an adaptation of Joint-Embedding Predictive Architectures (JEPA) to reinforcement learning from image observations. The authors use a vision transformer to encode three consecutive frames, a target encoder maintained as an exponential moving average, and a small predictor that receives the action; they train with a Euclidean latent prediction loss plus optional actor-critic gradients and a variance regularization loss. In a Cart Pole experiment with pixel observations and PPO, they compare four configurations (JEPA loss on/off, RL gradient propagation on/off, variance regularizer on/off) and report that the configuration with JEPA loss, no RL gradients, and variance regularization prevents collapse and learns. They conclude that JEPA can learn informative state representations without RL gradient propagation and are promising for image-based RL.
Significance. If the central claim were established, it would be a useful demonstration that a self-supervised latent prediction objective can support RL from images without direct reward-gradient flow through the encoder, potentially reducing sample complexity and improving transfer. The architecture and loss are clearly described, the collapse-prevention mechanism is concrete, and the three-way ablation is a sensible experimental strategy. However, the evidence is currently too thin to support the broad conclusion: one environment, five runs, no error bars, no comparison against standard representation baselines, and an ablation that does not isolate the JEPA loss from the variance regularizer. The contribution is more of a proof-of-concept than a validated method.
major comments (3)
- [Section 3, configuration 4; Eq. (2)] The central sub-claim in Section 3 ('This shows that JEPA is able to learn informative state representations without gradient propagation from the reinforcement learning task') rests on configuration 4 (J, no gradient propagation, R), but the ablation compares configuration 4 with configuration 3 (J, no gradient propagation, no R), which differs by two factors: the presence of the JEPA loss and the presence of the regularizer. Because Eq. (2) is present in configuration 4 and absent in configuration 3, the comparison cannot isolate the contribution of the JEPA loss from that of the variance regularizer. A control with no JEPA loss but with the same regularizer (J omitted, no gradient propagation, R) is required; without it, the observed learning could be driven entirely by the regularizer, which is known to prevent collapse and produce non-degenerate embeddings. This gap is load-bearing because the conclusion attributes the learning to JEPA rather than to variance regularization.
- [Section 3, Figure 3] The empirical evidence is too thin to support the paper's conclusions. Figure 3 shows only the averaged episodic return of five runs on Cart Pole, with no error bars, standard deviations, or seed-level analysis, and no quantitative comparison of final performance or sample efficiency. There is also no baseline against standard representation-learning methods for RL (e.g., a VAE or CURL), despite the introduction motivating JEPA as an alternative to autoencoders. The statement in Section 4 that 'JEPAs are promising candidates for reinforcement learning' therefore goes beyond what the presented data support.
- [Sections 1 and 4] The introduction states that the method targets 'problems like Atari games' with low-dimensional state vectors, but the only experiment is Cart Pole. Since Cart Pole has a four-dimensional state and very simple visual dynamics, it is not representative of the harder image-based RL tasks the method is proposed for. The conclusion should be explicitly restricted to the demonstrated setting unless additional environments and comparisons are provided.
minor comments (5)
- [Section 2, target encoder update] The formula 'θt+1 = 0.99·θt + 0.01·θt+1' is self-referential; the intended exponential moving average should use distinct notation for the target and online encoder weights.
- [Section 2, Figure 1] The introduction mentions an additional latent variable z in the JEPA formulation, but z never appears in the described method or in Eq. (1); clarify whether z is omitted and why.
- [Section 3, Figure 3] The caption says 'accumulated results' and the text says 'running average,' but neither the aggregation window nor the accumulation procedure is specified.
- [Section 2, Eq. (2)] The term Var(sx)_i is not defined precisely; specify whether this is the sample variance over the batch for each embedding dimension and whether the batch is the same as that used for the JEPA loss.
- [Section 2, reproducibility] The phrase 'a very limited number of na neurons' appears to contain a typo, and PPO hyperparameters, vision transformer architecture, image resolution, optimizer, and batch size are not reported, which limits reproducibility.
Circularity Check
No significant circularity: the paper's claims are empirical ablation results, not derivations that reduce to their inputs.
full rationale
The paper's core claims are empirical: it adapts a JEPA loss (Eq. 1), adds a variance regularizer (Eq. 2), and compares four configurations in Cart Pole. None of the target conclusions is defined in terms of the quantities it purports to demonstrate. The JEPA loss is a self-supervised latent matching objective, not a fitted predictor of the reported RL returns; no parameter is tuned to reproduce the empirical result. The comparison between configuration 3 (J, ∇hat, Rhat) and configuration 4 (J, ∇hat, R) is used to argue that the regularizer prevents collapse and that JEPA can learn informative representations without RL gradient propagation. This inference is vulnerable to a missing control: the paper never runs (Jhat, ∇hat, R), i.e., the variance regularizer alone, so one cannot fully attribute the learned behavior to the JEPA objective rather than to the regularizer. But that is an experimental-control and attribution issue, not circularity: it does not make the result equivalent to its inputs by construction. There are no load-bearing self-citations; the cited works (LeCun, BYOL, I-JEPA, VICReg) are external prior art, not the authors' own results. Accordingly, the appropriate circularity score is 0.
Assumptions & free parameters
free parameters (3)
- embedding dimension demb =
64
- EMA coefficient for target encoder =
0.99
- variance regularization clamp threshold =
1
assumptions (4)
- domain assumption The vision transformer with positional encoding over i, j, t captures the spatio-temporal information needed for state prediction.
- domain assumption The action taken between st and st+1, injected into the predictor, is the only extra information needed to predict the next latent embedding.
- domain assumption Batch-wise variance above 1 per embedding dimension is a sufficient proxy for informativeness and collapse prevention.
- domain assumption The Euclidean distance in latent space is an appropriate similarity measure for the prediction task.
Cite this review
Pith. "Pith review of JEPA for RL: Investigating Joint-Embedding Predictive Architectures for Reinforcement Learning." pith.science (2026). https://pith.science/paper/2ITX3VUD
@misc{pith2026250416591,
author = {Pith},
title = {Pith review of: JEPA for RL: Investigating Joint-Embedding Predictive Architectures for Reinforcement Learning},
year = {2026},
howpublished = {\url{https://pith.science/paper/2ITX3VUD}},
note = {Machine review of arXiv:2504.16591}
}
read the original abstract
Joint-Embedding Predictive Architectures (JEPA) have recently become popular as promising architectures for self-supervised learning. Vision transformers have been trained using JEPA to produce embeddings from images and videos, which have been shown to be highly suitable for downstream tasks like classification and segmentation. In this paper, we show how to adapt the JEPA architecture to reinforcement learning from images. We discuss model collapse, show how to prevent it, and provide exemplary data on the classical Cart Pole task.
Figures
Reference graph
Works this paper leans on
-
[1]
Human-level control through deep reinforcement learning
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):529–533, 2015
work page 2015
-
[2]
Gymnasium: A standard interface for reinforcement learning environments
Mark Towers and et al. Gymnasium: A standard interface for reinforcement learning environments. arXiv preprint arXiv:2407.17032 , 2024
arXiv 2024
-
[3]
Curl: Contrastive unsupervised representations for reinforcement learning, 2020
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning, 2020
work page 2020
-
[4]
Auto-encoding variational bayes
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. Proceedings of the 2nd International Conference on Learning Representations (ICLR) , 2014
work page 2014
-
[5]
David Ha and J¨ urgen Schmidhuber. World models. Neural Information Processing Sys- tems (NeurIPS) , 2018
work page 2018
-
[6]
Christopher P. Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in β-VAE. 4 2018
work page 2018
-
[7]
A path towards autonomous machine intelligence
Yann LeCun. A path towards autonomous machine intelligence. arXiv preprint arXiv:2205.12868, 2022
arXiv 2022
-
[8]
Bootstrap your own latent-a new approach to self-supervised learning
Jean-Bastien Grill, Florian Strub, Florent Altch´ e, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems , 33:21271–21284, 2020
work page 2020
Show all 10 references
-
[9]
Self-supervised learning from im- ages with a joint-embedding predictive architecture
Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from im- ages with a joint-embedding predictive architecture. In Proceedings of the IEEE/CVF Conference on Computer Vision and...
2023
-
[10]
Vicreg: Variance-invariance-covariance regularization for self-supervised learning
Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Variance-invariance-covariance regularization for self-supervised learning. CoRR, abs/2105.04906, 2021
2021 arXiv
Reviewed August 16, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.