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.

Teaching an AI agent to handle new tasks it has never seen before is a key challenge in robotics and machine learning. A standard approach trains agents on randomly generated reward functions. But we show this is flawed: in high dimensions, random tasks become nearly indistinguishable, causing the training signal to collapse. To fix this, we extract task descriptions directly from recorded experience data, building a distribution of tasks grounded in what the agent can actually do. This simple drop-in replacement requires no other changes to existing algorithms. Across three benchmarks and seven baseline methods, our approach improves zero-shot performance by an average of 20%.