JEPA · 2026-10-14
Sparse-TD-JEPA: Improving Zero-Shot Reinforcement Learning With Sparse Latent Task Representations
Duaa Nakshbandi, Zhiyu An, Wan Du
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
Sparse-TD-JEPA adds an ℓ1 sparsity penalty during training to concentrate reward-predictive structure into a smaller subspace of the task encoder’s latent space, improving zero-shot RL via TD-JEPA. The authors evaluate on 65 tasks across DMC and OGBench and report higher mean performance than TD-JEPA in all 13 domain-level comparisons, with analysis showing a sharper inferred reward landscape.
Why it matters
In successor-feature zero-shot RL, dense task encoders can yield diffuse inferred reward signals that reduce contrast between goal-relevant states and background, weakening the policy signal. Sparse-TD-JEPA targets this by encouraging per-example sparsity in the projected task embedding, while keeping the closed-form linear reward inference procedure unchanged.
Method
- Identify that dense task-encoder representations can produce diffuse inferred reward signals, flattening the predicted reward landscape and weakening the downstream policy’s signal.
- Add an ℓ1 sparsity penalty on the projected task encoder output (applied only during training) to concentrate reward-predictive information into a smaller effective subspace.
- Instantiate as Sparse-TD-JEPA with a single-line code addition to the TD-JEPA training update; no architectural change or modification to closed-form inference.
Limitation
The paper text does not state a specific limitation on method applicability or performance beyond the described mechanism and evaluations.
Abstract (from arXiv)
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. We identify an underexplored limitation of this approach: when the task encoder produces dense representations, the inferred reward signal becomes diffuse across the state space, reducing the contrast between goal-relevant states and the background and providing a weaker signal for the policy to follow. We address this limitation by introducing an l1 sparsity penalty that concentrates reward-predictive structure into a smaller effective subspace of the task encoder's latent space. The penalty is applied only during training and requires no architectural change or modification to the closed-form inference procedure. We instantiate this approach as Sparse-TD-JEPA, a single-line of code addition based on Temporal-Difference Joint-Embedding Predictive Architecture (TD-JEPA) that significantly improved its performance in zero-shot reinforcement learning. We evaluate Sparse-TD-JEPA on 65 tasks across the DeepMind Control (DMC) Suite and OGBench, covering locomotion, navigation, trajectory stitching, puzzle solving, and robotic manipulation. Sparse-TD-JEPA improves over TD-JEPA in all 13 domain-level comparisons. On DMC, it increases the average return from 661 +- 6 to 714 +- 3. On OGBench, it increases 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. Overall, Sparse-TD-JEPA achieves the highest mean performance among established zero-shot baselines in nine of the 13 domains. Analysis further shows that sparsity produces a sharper and more task-informative inferred reward landscape.
Related papers
- VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learning
- Disentangling Task-Relevant Latent Dynamics for Reward Modeling
- Vision-Predictor-Action: Towards Efficient Embodied AI via Neuro-Symbolic JEPA Predictors
- CausalDreamer: Learning Predictive World Models with Latent Disentanglement
- DiffusionART: Efficient, Stable and Scalable Reinforcement of Diffusion Models
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-09 · Indexed by Pub-AI: 2026-10-14 UTC
Sparse-TD-JEPA: Improving Zero-Shot Reinforcement Learning With Sparse Latent Task Representations
Preprint posted: 2026-10-09 · Original record ↗
- arxiv
- 2610.14400
Explicit source versions (1)
- v1 · 2026-10-09
- Preprint available
Source date: 2026-10-09 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
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. We identify an underexplored limitation of this approach: when the task encoder produces dense representations, the inferred reward signal becomes diffuse across the state space, reducing the contrast between goal-relevant states and the background and providing a weaker signal for the policy to follow. We address this limitation by introducing an l1 sparsity penalty that concentrates reward-predictive structure into a smaller effective subspace of the task encoder's latent space. The penalty is applied only during training and requires no architectural change or modification to the closed-form inference procedure. We instantiate this approach as Sparse-TD-JEPA, a single-line of code addition based on Temporal-Difference Joint-Embedding Predictive Architecture (TD-JEPA) that significantly improved its performance in zero-shot reinforcement learning. We evaluate Sparse-TD-JEPA on 65 tasks across the DeepMind Control (DMC) Suite and OGBench, covering locomotion, navigation, trajectory stitching, puzzle solving, and robotic manipulation. Sparse-TD-JEPA improves over TD-JEPA in all 13 domain-level comparisons. On DMC, it increases the average return from 661 +- 6 to 714 +- 3. On OGBench, it increases 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. Overall, Sparse-TD-JEPA achieves the highest mean performance among established zero-shot baselines in nine of the 13 domains. Analysis further shows that sparsity produces a sharper and more task-informative inferred reward landscape.