# Improving Zero-Shot Offline RL via Behavioral Task Sampling

arXiv: https://arxiv.org/abs/2604.25496 · HTML: https://mlsanity.com/paper/2604.25496/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 44 of its 56 references that are in the map.
- First posted on arXiv 28 Apr 2026.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 30 of 44)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Flamingo: a Visual Language Model for Few-Shot Learning](https://mlsanity.com/paper/2204.14198/) (2022)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Exploration by Random Network Distillation](https://mlsanity.com/paper/1810.12894/) (2018)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Learning Robust Rewards with Adversarial Inverse Reinforcement Learning](https://mlsanity.com/paper/1710.11248/) (2017)
- [Guided Cost Learning: Deep Inverse Optimal Control via Policy Optimization](https://mlsanity.com/paper/1603.00448/) (2016)
- [Successor Features for Transfer in Reinforcement Learning](https://mlsanity.com/paper/1606.05312/) (2016)
- [Language to Rewards for Robotic Skill Synthesis](https://mlsanity.com/paper/2306.08647/) (2023)
- [Automatic Goal Generation for Reinforcement Learning Agents](https://mlsanity.com/paper/1705.06366/) (2017)
- [IQ-Learn: Inverse soft-Q Learning for Imitation](https://mlsanity.com/paper/2106.12142/) (2021)
- [Skew-Fit: State-Covering Self-Supervised Reinforcement Learning](https://mlsanity.com/paper/1903.03698/) (2019)
- [Automatic Curriculum Learning For Deep RL: A Short Survey](https://mlsanity.com/paper/2003.04664/) (2020)
- [Learning One Representation to Optimize All Rewards](https://mlsanity.com/paper/2103.07945/) (2021)
- [Universal Successor Features Approximators](https://mlsanity.com/paper/1812.07626/) (2018)
- [Maximum Entropy Gain Exploration for Long Horizon Multi-goal Reinforcement Learning](https://mlsanity.com/paper/2007.02832/) (2020)
- [Don't Change the Algorithm, Change the Data: Exploratory Data for Offline Reinforcement Learning](https://mlsanity.com/paper/2201.13425/) (2022)
- [Learning with AMIGo: Adversarially Motivated Intrinsic Goals](https://mlsanity.com/paper/2006.12122/) (2020)
- [Autotelic Agents with Intrinsically Motivated Goal-Conditioned Reinforcement Learning: a Short Survey](https://mlsanity.com/paper/2012.09830/) (2020)
- [Does Zero-Shot Reinforcement Learning Exist?](https://mlsanity.com/paper/2209.14935/) (2022)
- [Open-Endedness is Essential for Artificial Superhuman Intelligence](https://mlsanity.com/paper/2406.04268/) (2024)
- [Zero-Shot Whole-Body Humanoid Control via Behavioral Foundation Models](https://mlsanity.com/paper/2504.11054/) (2025)

## Built on by (most cited first, 1 of 1)

- [Behavioral Foundation Models for Quality Diversity](https://mlsanity.com/paper/2609.35615/) (2026)

## Abstract (arXiv)

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.