Sparse-TD-JEPA: Improving Zero-Shot Reinforcement Learning With Sparse Latent Task Representations

The gist: Sparse-TD-JEPA introduces an l1 sparsity penalty applied to the task encoder's latent representations to concentrate reward-predictive structure into a smaller effective subspace, significantly improving zero-shot reinforcement learning performance across diverse benchmarks.

How it works

Successor-feature-based zero-shot reinforcement learning methods pre-train a task encoder and successor-feature predictor on offline, reward-free data, then recover a policy at test time by solving a closed-form linear regression over a small labeled dataset <ref:2610.14400#pg6>. The paper identifies an underexplored limitation where dense representations produce diffuse reward signals across the state space, reducing contrast between goal-relevant states and the background <ref:2610.14400#pg9>. To address this, Sparse-TD-JEPA introduces an l1 sparsity penalty applied only during training that concentrates reward-predictive structure into a smaller effective subspace of the task encoder’s latent space <ref:2610.14400#pg10>. This approach is implemented as a single-line code addition based on Temporal-Difference Joint-Embedding Predictive Architecture (TD-JEPA) <ref:2610.14400#pg10>.

The training objective for Sparse TD-JEPA is augmented with the sparsity penalty, defined as Lsp = 1/B sum i project z(ψ(si)) / 1 <ref:2610.14400#pg10>. The full training objective is then L = Lϕ + Lψ + λϕLorth(ϕ) + λψLorth(ψ) + λsp Lsp <ref:2610.14400#pg10>. This penalty is applied only during training and requires no architectural change or modification to the closed-form inference procedure <ref:2610.14400#pg10>.

Theoretical Analysis

The theoretical analysis characterizes what the penalty optimizes under the TD-JEPA sphere constraint, showing that it acts as an exact monotone measure of effective dimension <ref:2610.14400#pg10>. Specifically, Proposition 1 shows that l1 = p deff(x), meaning the penalty is a strictly increasing function of the effective dimension deff(x) <ref:2610.14400#pg10>. This implies that the penalty carries no information about representation magnitude because it is fixed by the constraint <ref:2610.14400#pg10>.

The analysis further shows that for p=1, l1 is the convex member of the family with the largest sparse-to-dense dynamic range on the sphere <ref:2610.14400#pg10>. Proposition 3 describes a gradient flow where mass is transported from coordinates below a state-dependent participation threshold τ(x) to those above it, meaning strongly active coordinates grow and weakly active ones decay <ref:2610.14400#pg10>.

Experimental Results

Sparse TD-JEPA improves over the dense baseline in almost every evaluated domain <ref:2610.14400#pg10>. On DeepMind Control (DMC), Sparse TD-JEPA increases the average return from 661±6 to 714±3, a gain of 53 points <ref:2610.14400#pg9>. On OGBench, it raises the average success rate from 38.0 ± 0.8 to 48.4 ± 2.3, with improvements of 13.4–18.0 percentage points on the evaluated antmaze navigation and stitching domains <ref:2610.14400#pg10>.

Analysis Experiment 1 demonstrates that the sparse encoder produces a sharper, more peaked predicted reward landscape than the dense encoder <ref:2610.14400#pg10>. This is verified by measuring the correlation between the Dense and Sparse model’s learned value estimate V(s) = ψ(s)⊤zr against a distance-based task progress signal x, showing a statistically significant higher correlation for the sparse encoder in all domains except for Large Navigate <ref:2610.14400#pg10>.

Failure Mode Explanation

The paper investigates why the dense baseline fails, observing that the predicted reward stays elevated across most of the goal room rather than peaking sharply at the goal itself, which gives little local gradient to follow <ref:2610.14400#pg10>. Sparse TD-JEPA resolves this by producing a more informative reward landscape where small differences in position are reflected in V(s) for the dense encoder, allowing the policy to commit to the goal <ref:2610.14400#pg10>. This is illustrated by Failure Mode 2, where the dense baseline’s V(s) fluctuates without a clear upward trend even as it moves in and out of the goal room <ref:2610.14400#pg10>.

Conclusion

Sparse TD-JEPA improves over the dense baseline in almost every evaluated domain <ref:2610.14400#pg10>. The gains are tied to a sharper, more goal-concentrated inferred reward that persists across labeled-set sizes <ref:2610.14400#pg9>. Future work includes applying this sparsity penalty to other successor-feature-based zero-shot RL methods beyond TD-JEPA <ref:2610.14400#pg10>. The paper concludes that the combination of per-example l1 concentration under a global second-moment isotropy constraint is the classical recipe of sparse dictionary learning <ref:2610.14400#pg10>.

AI Use Statement

In this work, we used AI tools (LLM-based assistants) to assist with polishing the manuscript text <ref:2610.14400#pg10>. We have not used generative AI tools for research ideation or method design <ref:2610.14400#pg10>. The identification of the failure mode, the proposed sparsity penalty, the theoretical analysis, and all experimental conclusions are the authors’ own <ref:2610.14400#pg10>. We have reviewed all AI-assisted work: AI-assisted text was checked against the underlying experimental results <ref:2610.14400#pg10>. We take responsibility for the final content of this work, including text, claims, and artifacts produced with the aid of generative AI <ref:2610.14400#pg10>.

Reproducibility Statement

We have taken several steps to make our results reproducible <ref:2610.14400#pg10>. All theoretical claims summarized in Section 4.2 are proved in Appendix A, which states every assumption and includes complete proofs <ref:2610.14400#pg10>. Appendix C documents the seeds and the result-aggregation protocol <ref:2610.14400#pg10>. Appendix D specifies all baselines and their hyperparameter grids <ref:2610.14400#pg10>. Appendix E gives full details of the datasets, environments, and evaluation protocol <ref:2610.14400#pg10>.

References

Mahmoud Assran et al. Self-supervised learning from images with a joint-embedding predictive architecture <ref:2610.14400#pg10>. Mido Assran et al. V-jepa 2: Self-supervised video models enable understanding, prediction and planning <ref:2610.14400#pg10>. Marco Bagatella et al. Tdjepa: Latent-predictive representations for zero-shot reinforcement learning <ref:2610.14400#pg10>. Adrien Bardes et al. Revisiting feature prediction for learning visual representations from video <ref:2610.14400#pg10>. Andre Barreto et al. Successor features for transfer in reinforcement learning <ref:2610.14400#pg10>. David L Donoho. Compressed sensing <ref:2610.14400#pg10>. Dibya Ghosh et al. Reinforcement learning from passive data via latent intentions <ref:2610.14400#pg10>. PMLR, 23 <ref:2610.14400#pg10>. Jean-Bastien Grill et al. Bootstrap your own latent-a new approach to self-supervised learning <ref:2610.14400#pg10>. Daniel Lawson et al. Selfpredictive representations for combinatorial generalization in behavioral cloning <ref:2610.14400#pg10>. Yann LeCun et al. A path towards autonomous machine intelligence version 0.9. 2, 2, 2, 2-6-27 <ref:2610.14400#pg10>. Vincent Liu et al. The utility of sparse representations for control in reinforcement learning <ref:2610.14400#pg10>. Alireza Makhzani and Brendan Frey. K-sparse autoencoders <ref:2610.14400#pg10>. Bruno A Olshausen and David J Field. Emergence of simple-cell receptive field properties by learning a sparse code for natural images <ref:2610.14400#pg10>. Seohong Park et al.

1. The Sparse-TD-JEPA system can be improved by incorporating a mechanism to explicitly identify and enforce shared low-dimensional structures among task embeddings, as suggested by the analysis: We note that Eq. equation 5 constrains each embedding individually and does not by itself force different states to share a support; whether a shared low-dimensional structure emerges is an empirical question.

2. The system can be enhanced to leverage the theoretical findings regarding the penalty's optimization dynamics: "Sparse TD-JEPA this pressure acts jointly with Lϕ, Lψ, and the orthonormality losses of Eq. equation 6, so training equilibrates at an intermediate effective dimension determined by the trade-off between predictive fit and the sparsity coefficient λsp."

3. The improved system can be designed to specifically exploit the relationship between reward landscape sharpness and task progress: Sparse encoder produced a statistically significant higher correlation between predicted reward and task completion distance than Dense encoder in all domains except for Large Navigate, allowing the policy to have a sharper, more concentrated inferred reward landscape with a stronger signal for the downstream policy to follow.

4. The system can be refined to better handle limited labeled data during inference by recognizing that the gap between sparse and dense encoders is persistent: "the sparse encoder’s advantage in how well V (s) tracks the ground-truth signal is not merely a low-data artifact tied to the N = 104 inference budget used in the main experiments, but persists as more labeled data becomes available."