Zero-shot RL stops losing the goal in a dense task code: one sparsity penalty lifts DMC average return by 53 points

Zero-shot reinforcement learning pre-trains a task encoder and recovers a policy at test time by solving a closed-form linear regression over a small labeled dataset; when the encoder output is dense, the inferred reward spreads across the state space, weakening the contrast between goal-relevant states and the background and leaving the policy a weaker signal to follow. Earlier successor-feature and task-encoder methods focused on how to construct the task representation, for example through transition-graph eigenfunctions, intention-conditioned values, or temporal distances, while the bottleneck caused by a dense encoder output received comparatively little attention. Across 65 tasks on DeepMind Control and OGBench, the sparse version beats the dense TD-JEPA baseline in all 13 domain-level comparisons; DMC average return rises by 53 points and OGBench average success rate by 10.4 percentage points, with 13.4 to 18.0 percentage point gains on the evaluated antmaze navigation and stitching domains.

The method applies a sparsity penalty to the projected task embeddings during training, so each embedding concentrates its mass on fewer strongly active dimensions and the reward-predictive structure occupies a smaller effective subspace. The penalty acts only during training, requires no architectural change, and leaves the closed-form inference procedure untouched, so the implementation is a single added line in the TD-JEPA update step. Evaluation covers 20 tasks across the DMC walker, cheetah, quadruped, and pointmass domains, plus 45 tasks across nine OGBench domains, with 10 independent rollouts per task averaged.

Analysis shows the sparse encoder produces a sharper inferred reward landscape that peaks near the goal, offering a mechanistic explanation for the zero-shot performance gains. With the dense encoder the inferred reward stays elevated across most of the goal room and gives little indication of the exact goal location once the agent is inside; with the sparse encoder the same signal peaks sharply at the goal and falls off toward the room's edges. Over 20 independently resampled observation batches, the correlation between inferred reward and a task-progress signal is measured with 5-fold cross-validation; the sparse encoder is significantly higher in six of nine domains, and the gap stays roughly flat as the labeled-set size sweeps across several orders of magnitude.

The result applies to zero-shot reinforcement learning built on successor features and task encoders: pre-training on reward-free offline data, then recovering a policy at test time by closed-form linear regression over a small labeled dataset. Evaluation covers 20 tasks across four DMC domains and 45 tasks across nine OGBench domains, spanning locomotion, navigation, trajectory stitching, puzzle solving, and robotic manipulation. The sparsity penalty is instantiated only on TD-JEPA; the authors propose applying it to other successor-feature-based zero-shot methods and using the resulting low-effective-dimension task embeddings to understand what individual latent dimensions encode.

The penalty constrains each embedding individually and does not force different states to share a support, so whether a shared low-dimensional structure emerges is an empirical question. In the distance-reward correlation analysis the Large Navigate domain does not reach statistical significance, and in the labeled-set-size sweep the puzzle-3x3 domain shows no reliable signal for either model. Whether the sparse encoder better represents task dimensions beyond distance-to-goal is left to future work.