Improving Zero-Shot Offline RL via Behavioral Task Sampling

1. Problem Statement

Offline Zero-Shot Reinforcement Learning (ZSRL) aims to train agents that can optimize unseen reward functions without further interaction with the environment, relying solely on a pre-collected offline dataset.

• Standard Approach: Current state-of-the-art methods (e.g., Successor Features, Forward-Backward representations) learn a state embedding $\phi(s)$ and train a policy conditioned on a task vector $z$. The reward function is defined linearly as $R_z(s) = \phi(s)^\top z$.
• The Flaw: To train these policies, existing algorithms typically sample task vectors $z$ uniformly from a high-dimensional unit hypersphere ($S^{d-1}$).
• The Core Issue: The authors argue that this uniform sampling creates a mismatch between the geometry of the task space and the physics of the environment. In high-dimensional spaces, the "concentration of measure" phenomenon causes uniformly sampled vectors to be nearly orthogonal to the actual behaviors the agent can produce (the behavioral space). This leads to "signal dilution," where the reward signal becomes vanishingly small for almost all behaviors, making it impossible for the agent to distinguish optimal actions from suboptimal ones. This results in poor zero-shot generalization.

2. Methodology: Behavioral Task Distribution (BTD)

The authors propose replacing uniform task sampling with a data-driven approach called Behavioral Task Sampling.

A. Theoretical Insight

• Feature Occupancy: They define the feature occupancy $\psi_\pi$ of a policy as the discounted sum of state features along its trajectory.
• Vanishing Variance: They prove (Proposition 4.1) that as the latent dimension $d$ increases, the expected variance of returns across policies under uniform task sampling converges to zero. This means the learning signal collapses, preventing effective policy optimization.

B. The Proposed Algorithm

Instead of sampling $z$ from a uniform distribution, the method extracts task vectors directly from the offline dataset:

1. Extraction: From the offline dataset $D$, the authors extract sub-trajectories $\tau$. For each sub-trajectory, they compute the empirical feature occupancy $\tilde{\psi}_\tau = \sum \gamma^t \phi(s_t)$.
2. Task Vector Definition: The task vector for a trajectory is defined as the normalized occupancy: $z_\tau = \tilde{\psi}_\tau / \|\tilde{\psi}_\tau\|_2$. This vector represents the direction in feature space that maximizes the reward for that specific observed behavior.
3. Distribution Modeling: The set of all extracted vectors forms an empirical distribution $p_{data}$. The authors fit a parametric density model (specifically a Gaussian Mixture Model or GMM) to this data, creating the Behavioral Task Distribution (BTD), denoted as $p_\theta(z)$.
4. Training: During policy training, task vectors are sampled from the learned BTD ($z \sim p_\theta(z)$) rather than the unit hypersphere. This ensures the agent is trained on tasks that are physically achievable and grounded in the environment's dynamics.

3. Key Contributions

1. Identification of Signal Dilution: The paper theoretically and empirically demonstrates that uniform task sampling in high-dimensional spaces leads to a collapse of the learning signal due to orthogonality with the behavioral space.
2. Behavioral Task Distribution (BTD): A novel, method-agnostic framework that extracts and models the distribution of achievable tasks directly from offline data.
3. Integration: The method is a drop-in replacement for the task sampling step in existing ZSRL algorithms (Successor Features and Forward-Backward) without modifying the underlying representation learning objectives.
4. Comprehensive Evaluation: Extensive experiments across multiple benchmarks, representation learning methods, and latent dimensions.

4. Experimental Results

The authors evaluated their method on standard ZSRL benchmarks (Cheetah, Walker, Quadruped) using datasets collected via Random Network Distillation (RND). They compared against seven baselines (including Autoencoders, Transition Models, LRA, and BYOL) across various latent dimensions.

• Performance Gain: BTD-sampling improved zero-shot performance by an average of 20% across all environments and baselines.
• Robustness to Dimensionality:
  • Baseline Failure: As the latent dimension increased (up to 1000), standard uniform sampling methods suffered a sharp performance drop, often collapsing to near-random performance.
  • BTD Stability: BTD-sampling maintained high performance even at high dimensions (e.g., 1000), proving it effectively mitigates the signal dilution problem.
• Reduced Variance: BTD-sampling significantly reduced the variance in test performance across different representation learning methods, indicating more consistent and reliable learning.
• Ablation Studies:
  • Mixing Strategies: Blending uniform sampling with BTD sampling (interpolating between the two) degraded performance, confirming that introducing arbitrary, unachievable tasks harms the learning signal.
  • GMM Components: Performance plateaued after ~20 GMM components, suggesting the behavioral task space can be captured with a relatively simple distribution.
  • Visualization: UMAP projections showed that uniform sampling covers the entire hypersphere, while the actual behavioral tasks occupy a small, specific manifold, validating the need for data-driven sampling.

5. Significance and Conclusion

This work fundamentally shifts the paradigm of offline ZSRL from arbitrary task generation to behavioral grounding.

• Practical Impact: It highlights that in offline settings, what tasks are chosen for training is as critical as how the policy learns them.
• Scalability: By decoupling the task distribution from the curse of dimensionality inherent in uniform sampling, the method enables the training of generalist agents in high-dimensional latent spaces.
• Generalizability: The approach is compatible with various representation learning techniques (SF, FB, Autoencoders, etc.), making it a scalable path toward more robust and generalizable reinforcement learning agents for real-world deployment where tasks are physically constrained.

In summary, the paper argues that to achieve true zero-shot generalization in offline RL, agents must be trained on tasks that reflect the actual capabilities and dynamics of the environment, rather than an abstract, uniform distribution of mathematical possibilities.