Improving Zero-Shot Offline RL via Behavioral Task Sampling

• URL: http://arxiv.org/abs/2604.25496v1
• Date: Tue, 28 Apr 2026 10:56:54 GMT
• Title: Improving Zero-Shot Offline RL via Behavioral Task Sampling
• Abstract summary: offline zero-shot reinforcement learning aims to learn agents that optimize unseen reward functions without additional environment interaction.<n>In most existing algorithms, task vectors are randomly sampled, implicitly assuming this adequately captures the structure of the task space.<n>We argue that doing so leads to suboptimal zero-shot generalization.<n>We introduce a simple and general reward function extraction procedure that integrates into existing offline zero-shot RL algorithms.
• Score: 5.045010561958866
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Beyond Solvability: Task Learnability as a Static Prior for LLM RL Post-Training [52.78248352851176]
  Reinforcement learning (RL) has become a central post-training paradigm for eliciting reasoning capabilities in large language models.<n>We study this residual axis as task learnability: a regime-conditional measure of expected positive response to continued training.<n>We propose TrajVal, a lightweight probe-based estimator that approximates per-task learnability from a short probe run and two endpoint evaluations.
  arXiv  Detail & Related papers  (2026-08-10T07:43:05Z)
• Soft Forward-Backward Representations for Zero-shot Reinforcement Learning with General Utilities [55.52456333246921]
  Forward-backward algorithms can retrieve a family of policies that can approximately solve any standard RL problem.<n>We tackle the greater problem class of RL with general utilities, in which the objective is an arbitrary differentiable function of the occupancy measure.<n>We show this additional complexity can be captured by a novel, maximum entropy (soft) variant of the forward-backward algorithm.
  arXiv  Detail & Related papers  (2026-02-06T15:24:37Z)
• Zero-Shot Off-Policy Learning [9.729890516322781]
  Off-policy learning methods seek to derive an optimal policy directly from a fixed dataset of prior interactions.<n>In this work, we address the off-policy problem in a zero-shot setting by discovering a theoretical connection of successor measures to stationary density ratios.<n>Our algorithm can infer optimal importance sampling ratios, effectively performing a stationary distribution correction with an optimal policy for any task on the fly.
  arXiv  Detail & Related papers  (2026-02-02T11:06:31Z)
• Fast Adaptation with Behavioral Foundation Models [82.34700481726951]
  Unsupervised zero-shot reinforcement learning has emerged as a powerful paradigm for pretraining behavioral foundation models.<n>Despite promising results, zero-shot policies are often suboptimal due to errors induced by the unsupervised training process.<n>We propose fast adaptation strategies that search in the low-dimensional task-embedding space of the pre-trained BFM to rapidly improve the performance of its zero-shot policies.
  arXiv  Detail & Related papers  (2025-04-10T16:14:17Z)
• Adaptive Reward Design for Reinforcement Learning [2.3031174164121127]
  We propose a suite of reward functions that incentivize an RL agent to complete a task specified by a formula as much as possible.<n>We develop an adaptive reward shaping approach that dynamically updates reward functions during the learning process.
  arXiv  Detail & Related papers  (2024-12-14T18:04:18Z)
• Unsupervised Zero-Shot Reinforcement Learning via Functional Reward Encodings [107.1837163643886]
  We present a functional reward encoding (FRE) as a general, scalable solution to this zero-shot RL problem. Our main idea is to learn functional representations of any arbitrary tasks by encoding their state-reward samples. We empirically show that FRE agents trained on diverse random unsupervised reward functions can generalize to solve novel tasks.
  arXiv  Detail & Related papers  (2024-02-27T01:59:02Z)
• Foundation Policies with Hilbert Representations [54.44869979017766]
  We propose an unsupervised framework to pre-train generalist policies from unlabeled offline data. Our key insight is to learn a structured representation that preserves the temporal structure of the underlying environment. Our experiments show that our unsupervised policies can solve goal-conditioned and general RL tasks in a zero-shot fashion.
  arXiv  Detail & Related papers  (2024-02-23T19:09:10Z)
• Zero-Shot Reinforcement Learning from Low Quality Data [5.079602839359521]
  Zero-shot reinforcement learning (RL) promises to provide agents that can perform any task in an environment after an offline, reward-free pre-training phase. Here, we explore how the performance of zero-shot RL methods degrades when trained on small homogeneous datasets. We propose fixes inspired by conservatism, a well-established feature of performant single-task offline RL algorithms.
  arXiv  Detail & Related papers  (2023-09-26T18:20:20Z)
• Hypernetworks for Zero-shot Transfer in Reinforcement Learning [21.994654567458017]
  Hypernetworks are trained to generate behaviors across a range of unseen task conditions. This work relates to meta RL, contextual RL, and transfer learning. Our method demonstrates significant improvements over baselines from multitask and meta RL approaches.
  arXiv  Detail & Related papers  (2022-11-28T15:48:35Z)
• Text Generation with Efficient (Soft) Q-Learning [91.47743595382758]
  Reinforcement learning (RL) offers a more flexible solution by allowing users to plug in arbitrary task metrics as reward. We introduce a new RL formulation for text generation from the soft Q-learning perspective. We apply the approach to a wide range of tasks, including learning from noisy/negative examples, adversarial attacks, and prompt generation.
  arXiv  Detail & Related papers  (2021-06-14T18:48:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.