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

Duaa Nakshbandi; Zhiyu An; Wan Du

cs.LGOct 9, 2026

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

Authors: Duaa Nakshbandi, Zhiyu An, Wan Du

Organizations: University of California, Merced

Abstract

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.

Figures & tables

Figure 1: We test our method for zero-shot performance on various control, locomotive, and exploration tasks.

	Laplacian	ICVF	HILP	FB	RLDP	BYOL	BYOL- $\gamma$	TD-JEPA	Ours
DMC (avg)	$591\pm 11$	$619\pm 10$	$620\pm 8$	$648\pm 4$	$610\pm 14$	$619\pm 11$	$645\pm 11$	$661\pm 6$	$\mathbf{714\pm 3}$
walker	$770\pm 5$	$727\pm 16$	$796\pm 8$	$\mathbf{812\pm 6}$	$724\pm 18$	$747\pm 11$	$786\pm 10$	$785\pm 7$	$\mathbf{814\pm 8}$
cheetah	$615\pm 19$	$606\pm 17$	$618\pm 6$	$673\pm 5$	$576\pm 45$	$623\pm 24$	$647\pm 9$	$689\pm 7$	$\mathbf{720\pm 6}$
quadruped	$635\pm 39$	$709\pm 14$	$695\pm 11$	$596\pm 9$	$665\pm 14$	$612\pm 28$	$683\pm 26$	$691\pm 5$	$\mathbf{801\pm 10}$
pointmass	$345\pm 22$	$436\pm 11$	$371\pm 37$	$\mathbf{513\pm 20}$	$\mathbf{476\pm 39}$	$\mathbf{493\pm 41}$	$465\pm 18$	$\mathbf{479\pm 24}$	$\mathbf{523\pm 40}$
OGBench (avg)	$14.8\pm 1.3$	$30.9\pm 0.6$	$38.0\pm 1.1$	$39.0\pm 0.7$	$27.1\pm 0.8$	$26.4\pm 0.8$	$30.4\pm 0.9$	$38.0\pm 0.8$	$\mathbf{48.4\pm 2.3}$

Table 1: Performance of zero-shot algorithms for DMC (reward) and OGBench (success rate) with proprioception inputs. We report means and standard errors across seeds; DMC values are rounded to the nearest integer and OGBench values to one decimal place. Numbers are bold for top algorithms if confidence intervals overlap.

Figure 2: Spatial reward landscape for task 3 on antmaze-large-navigate ( left pair ) and antmaze-medium-navigate ( right pair ). Colour encodes $z_{r}^{\top}\psi(s)$ normalized by the model’s own mean predicted reward for far-field states ( $>$ 10 m from goal). Within each domain, Dense TD-JEPA ( left ) and Sparse TD-JEPA ( right ) share the same colour scale. The star marks the goal; rings indicate 1 m, 3 m, and 5 m distance.

Figure 3: Correlation between each model’s learned value estimate

V(s)=\psi(s)\cdot z_{r}

and distance signal.

Figure 4: Example of linear regression correlation of Dense (Baseline) and Sparse task encoders.

Figure 5: Cross-validated

R^{2}

between learned value estimate and the distance signal.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Regularizer	Behavior on $\mathcal{S}_{d}$	Effect on training	Shown in
$\ell_{2}$ , $\\|x\\|_{2}^{2}$ , any $h(\\|x\\|_{2})$	constant	none (zero gradient)	Thm. 1 , Cor. 2
isotropically smoothed norms	constant	none (zero gradient)	Cor. 2
weight decay (output layer)	function-preserving	parameter scale only	Rem. 2
dropout on $\psi$ output	$\approx$ quadratic, $\to$ isotropic	noise; no sparsity pressure	Rem. 3
orthonormality (second moments)	orbit-invariant	blind to per-example $d_{\mathrm{eff}}(x)$	Lem. 2
$\ell_{p}$ , $2<p\leq\infty$	minimized at flat $x$	favors dense embeddings	Prop. 2 (iii)

Appendix

Table 2: Behavior of embedding-level regularizers under TD-JEPA’s sphere constraint.

d_{\mathrm{eff}}(x)

denotes the effective dimension of Definition 1 .

Figure 6: Full antmaze-large-stitch episode (task 4). Sparse TD-JEPA navigates from its start position to the goal, with $V(s)$ rising and distance-to-goal falling together over the episode. The dense TD-JEPA baseline instead loops within a small region far from the goal for the entire episode, never making net progress, and fails.

Figure 7: Full antmaze-medium-navigate episode (task 4). Sparse TD-JEPA reaches the goal directly from its start position. The dense TD-JEPA baseline remains tangled in a small loop immediately around its start position for the full episode and fails to make progress toward the goal.

Figure 8: Last five metres of an antmaze-large-stitch episode (task 4). Rollout zoomed to the 10 m $\times$ 10 m region centred on the goal. Sparse TD-JEPA’s value estimate $V(s)$ rises monotonically as it closes the remaining distance and succeeds; the dense TD-JEPA baseline’s $V(s)$ is noisy and directionless over the same last-mile window, even dropping sharply to near zero at one point, and the agent wanders near the goal room without ever committing to the goal, ultimately failing the episode.

Figure 9: Sparse-TD-JEPA achieves better zero-shot performance than the baseline framework in all tested domains, and outperforms all other established zero-shot baselines in most domains.

Domain group	lr ${}_{\text{features}}$	ortho. coef.
OGBench Antmaze	$\{10^{-4},10^{-5}\}$	$\{0,\,1\}$
OGBench Manipulation	$\{10^{-4},10^{-5}\}$	$\{0,\,1\}$
DMC	$\{10^{-4},10^{-5}\}$	$\{0,\,0.1,\,1\}$

Appendix

Table 3: Laplacian hyperparameter grid.

Domain group	lr ${}_{\text{features}}$	goal-sampling prob.
OGBench Antmaze	$\{10^{-4},10^{-5}\}$	$\{0.375,\,0.5\}$
OGBench Manipulation	$\{10^{-4},10^{-5}\}$	$\{0.375,\,0.5\}$
DMC	$\{10^{-4},10^{-5}\}$	$\{0.375,\,0.2,\,0.5\}$

Appendix

Table 4: HILP hyperparameter grid (also used, unmodified, for ICVF).

Domain group	lr b	ortho. coefficient
OGBench Antmaze	$\{10^{-4},10^{-5}\}$	$\{100,\,1000\}$
OGBench Manipulation	$\{10^{-4},10^{-5}\}$	$\{100,\,1000\}$
DMC	$\{10^{-4},10^{-5}\}$	$\{0.1,\,1,\,10\}$

Appendix

Table 5: FB hyperparameter grid.

Domain group	lr b	ortho. coefficient
OGBench Antmaze	$\{10^{-4},10^{-5}\}$	$\{0.001,\,0.01\}$
OGBench Manipulation	$\{10^{-4},10^{-5}\}$	$\{0.01,\,0.1\}$
DMC	$\{10^{-4},10^{-5}\}$	$\{0.001,\,0.01,\,0.1\}$

Appendix

Table 6: RLDP hyperparameter grid.

Domain group	lr ${}_{\text{features}}$	ortho. coefficient
OGBench Antmaze	$\{10^{-4},10^{-5}\}$	$\{0.001,\,0.01\}$
OGBench Manipulation	$\{10^{-4},10^{-5}\}$	$\{0.01,\,0.1\}$
DMC	$\{10^{-4},10^{-5}\}$	$\{0.001,\,0.01,\,0.1\}$

Appendix

Table 7: BYOL hyperparameter grid (also used, unmodified, for BYOL-

\gamma

Domain group	lr ψ	$\psi_{\text{ortho}}$ -coefficient
OGBench Antmaze	$\{10^{-4},10^{-5}\}$	$\{0.1,\,1\}$
OGBench Manipulation	$\{10^{-4},10^{-5}\}$	$\{1,\,10\}$
DMC	$\{10^{-4},10^{-5}\}$	$\{0.01,\,0.1,\,1\}$

Appendix

Table 8: TD-JEPA hyperparameter grid.

Domain	State dim.	Action dim.	Tasks	Task Description
Walker	24	6	Stand	Upright torso, minimal height
			Walk	Stand reward plus forward velocity
			Run	Stand reward plus higher velocity
			Spin	Continuous rotation while balanced
Cheetah	17	6	Walk	Reward linear in forward velocity, max 10 m/s
			Walk Backward	Reward linear in backward velocity

Appendix

Table 9: DMC dataset domains, dimensions, and tasks.

Domain	State dim	Action dim	Task	Task Description
Antmaze Medium Navigate	29	8	Task1–5	Navigate to five goals in the medium maze
Antmaze Large Navigate	29	8	Task1–5	Navigate to five goals in the large maze
Antmaze Medium Stitch	29	8	Task1–5	Reach five goals by stitching sub-trajectories
Antmaze Large Stitch	29	8	Task1–5	Reach five goals by stitching sub-trajectories
Antmaze Medium Explore	29	8	Task1–5	Reach five goals from low-quality, high-coverage exploratory data
Cube Single	28	5	Horizontal	Pick-and-place, horizontal offset

Appendix

Table 10: OGBench dataset domains, dimensions, and tasks.

Zero-Shot RL

See all 17 papers

Apr 28, 2026

Improving Zero-Shot Offline RL via Behavioral Task Sampling

Offline zero-shot reinforcement learning (RL) aims to learn agents that optimize unseen reward functions without additional environment interaction. The standard approach to this problem trains task-conditioned policies by sampling task vectors that define linear reward functions over learned state representations. In most existing algorithms, these task vectors are randomly sampled, implicitly assuming this adequately captures the structure of the task space. We argue that doing so leads to suboptimal zero-shot generalization. To address this limitation, we propose extracting task vectors directly from the offline dataset and using them to define the task distribution used for policy training. We introduce a simple and general reward function extraction procedure that integrates into existing offline zero-shot RL algorithms. Across multiple benchmark environments and baselines, our approach improves zero-shot performance by an average of 20%, highlighting the importance of principled task sampling in offline zero-shot RL.

Nazim Bendib, Nicolas Perrin-Gilbert, Olivier Sigaud

Institut des Systèmes Intelligents et de Robotique, ISIR, Sorbonne Université, Paris, France.

Zero-Shot RL Offline RL Multi-Task RL Reinforcement Learning Generalization in RL Goal-Conditioned RL+1

Ask PDF

Mar 10, 2026

Reward-Zero: Language Embedding Driven Implicit Reward Mechanisms for Reinforcement Learning

We introduce Reward-Zero, a general-purpose implicit reward mechanism that transforms natural-language task descriptions into dense, semantically grounded progress signals for reinforcement learning (RL). Reward-Zero serves as a simple yet sophisticated universal reward function that leverages language embeddings for efficient RL training. By comparing the embedding of a task specification with embeddings derived from an agent's interaction experience, Reward-Zero produces a continuous, semantically aligned sense-of-completion signal. This reward supplements sparse or delayed environmental feedback without requiring task-specific engineering. When integrated into standard RL frameworks, it accelerates exploration, stabilizes training, and enhances generalization across diverse tasks. Empirically, agents trained with Reward-Zero converge faster and achieve higher final success rates than conventional methods such as PPO with common reward-shaping baselines, successfully solving tasks that hand-designed rewards could not in some complex tasks. In addition, we develop a mini benchmark for the evaluation of completion sense during task execution via language embeddings. These results highlight the promise of language-driven implicit reward functions as a practical path toward more sample-efficient, generalizable, and scalable RL for embodied agents. Code will be released after peer review.

Heng Zhang, Haddy Alchaer, Arash Ajoudani +1

Human-Robot Interfaces and Interaction Lab, Istituto Italiano di Tecnologia, Genova, Italy · Ph.D. program of national interest in Robotics and Intelligent Machines (DRIM) and Universit ‘a di Genova, Genoa, Italy. · Edwardson School of Industrial Engineering, Purdue University, West Lafayette, IN, USA. +2

Reward Design for RL Dense Reward RL Reward Shaping Sparse-Reward RL

Ask PDF

Oct 23, 2025

A Unified Framework for Zero-Shot Reinforcement Learning

Zero-shot reinforcement learning (RL) has emerged as a setting for developing general agents, capable of solving downstream tasks without additional training or planning at test-time. While conventional RL optimizes policies for fixed rewards, zero-shot RL requires learning representations that enable immediate adaptation to arbitrary reward functions. As the field matures, the growing diversity of approaches demands a foundational framework reconciling different perspectives under a common unifying structure. In this work, we introduce a formal, unified framework for zero-shot RL, allowing for rigorous comparisons across methods. We propose a taxonomy organizing the algorithmic landscape along two levels: representation, distinguishing between compositional and direct methods based on their exploitation of action-value function decompositions; and learning paradigm, differentiating between reward-free and pseudo reward-free training. Additionally, we propose a unified view of existing error bounds, decomposing the total error into three primary contributing components: inference, reward, and approximation, serving as a foundation for more grounded comparisons of zero-shot methods.

Jacopo Di Ventura, Jan Felix Kleuker, Aske Plaat +1

Leiden University

Zero-Shot RL Formal Methods for RL Reward-Free RL Representation Learning for RL

Ask PDF

Abstract

Zero-Shot RL

Explore similar work

Improving Zero-Shot Offline RL via Behavioral Task Sampling

Reward-Zero: Language Embedding Driven Implicit Reward Mechanisms for Reinforcement Learning

A Unified Framework for Zero-Shot Reinforcement Learning