Table of Contents
Fetching ...

Zero-Shot Off-Policy Learning

Arip Asadulaev, Maksim Bobrin, Salem Lahlou, Dmitry Dylov, Fakhri Karray, Martin Takac

TL;DR

This work tackles zero-shot off-policy learning for Behavioral Foundation Models by tying Forward-Backward successor representations to stationary density ratios, enabling a training-free, occupancy-aware adaptation at test time. The proposed Zero-Shot Off-policy Learning (ZOL) computes a low-rank, FB-derived occupancy ratio surrogate $w_{\pi/\beta}(s,a)\approx (1-\gamma)\,W_{\pi_z}^\top B(s,a)$ and optimizes a latent task vector $z$ to maximize a corrected return while regularizing for stability and staying within dataset support. It infers an initial task embedding from reward information and then performs gradient-based latent search with rewards centered and weights normalized, avoiding extrapolation pitfalls typical of offline-to-online transfers. Across ExORL, OGBench, and SMPL Humanoid benchmarks, ZOL delivers robust improvements over strong zero-shot baselines, demonstrating the practical value of integrating occupancy-corrected latent adaptation with pretrained FB representations for fast, training-free task generalization.

Abstract

Off-policy learning methods seek to derive an optimal policy directly from a fixed dataset of prior interactions. This objective presents significant challenges, primarily due to the inherent distributional shift and value function overestimation bias. These issues become even more noticeable in zero-shot reinforcement learning, where an agent trained on reward-free data must adapt to new tasks at test time without additional training. 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. Using this insight, our algorithm can infer optimal importance sampling ratios, effectively performing a stationary distribution correction with an optimal policy for any task on the fly. We benchmark our method in motion tracking tasks on SMPL Humanoid, continuous control on ExoRL, and for the long-horizon OGBench tasks. Our technique seamlessly integrates into forward-backward representation frameworks and enables fast-adaptation to new tasks in a training-free regime. More broadly, this work bridges off-policy learning and zero-shot adaptation, offering benefits to both research areas.

Zero-Shot Off-Policy Learning

TL;DR

This work tackles zero-shot off-policy learning for Behavioral Foundation Models by tying Forward-Backward successor representations to stationary density ratios, enabling a training-free, occupancy-aware adaptation at test time. The proposed Zero-Shot Off-policy Learning (ZOL) computes a low-rank, FB-derived occupancy ratio surrogate and optimizes a latent task vector to maximize a corrected return while regularizing for stability and staying within dataset support. It infers an initial task embedding from reward information and then performs gradient-based latent search with rewards centered and weights normalized, avoiding extrapolation pitfalls typical of offline-to-online transfers. Across ExORL, OGBench, and SMPL Humanoid benchmarks, ZOL delivers robust improvements over strong zero-shot baselines, demonstrating the practical value of integrating occupancy-corrected latent adaptation with pretrained FB representations for fast, training-free task generalization.

Abstract

Off-policy learning methods seek to derive an optimal policy directly from a fixed dataset of prior interactions. This objective presents significant challenges, primarily due to the inherent distributional shift and value function overestimation bias. These issues become even more noticeable in zero-shot reinforcement learning, where an agent trained on reward-free data must adapt to new tasks at test time without additional training. 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. Using this insight, our algorithm can infer optimal importance sampling ratios, effectively performing a stationary distribution correction with an optimal policy for any task on the fly. We benchmark our method in motion tracking tasks on SMPL Humanoid, continuous control on ExoRL, and for the long-horizon OGBench tasks. Our technique seamlessly integrates into forward-backward representation frameworks and enables fast-adaptation to new tasks in a training-free regime. More broadly, this work bridges off-policy learning and zero-shot adaptation, offering benefits to both research areas.
Paper Structure (40 sections, 3 theorems, 47 equations, 8 figures, 7 tables)

This paper contains 40 sections, 3 theorems, 47 equations, 8 figures, 7 tables.

Key Result

Theorem 3.1

We write $w_{\pi}$ as an shorthand of $w_{\pi/\beta}$. Assume $d_\beta(s)>0$ whenever $d_\pi(s)>0$ and define the density ratio Then the centered reweighted objective satisfies Hence $\arg\max_\pi J_c(\pi) = \arg\max_\pi \mathbb{E}_{d_\pi}[r(s)]$.

Figures (8)

  • Figure 1: A sketch of the zero-shot off-policy adaptation. The agent completes the new task by combining actions from its experience.
  • Figure 2: ZOL improvements over other BFMs on ExORL (DMC-state) benchmark. For full results, refer to Tbl. \ref{['table:exorl_merged']}. Despite not having access to the environment, ZOL performs better, especially on long horizon tasks.
  • Figure 3: HumEnv benchmark. Performance improvement of ZOL over FB-CPR on SMPL Humanoid, regularized to diverse realistic motion capture trajectories from AMASS dataset.
  • Figure 4: Evaluation Benchmarks. We evaluate ZOL on seven robotic locomotion and manipulation environments across $5$ different tasks in unsupervised regime, which vary in complexity, data coverage, and reward types.
  • Figure 5: The offline dataset support (left), and the corresponding heatmaps of the reconstructed rewards from latents obtained by ZOL (third column) or FB (fourth column). Colors represent rewards (higher $\rightarrow$ brighter).
  • ...and 3 more figures

Theorems & Definitions (6)

  • Theorem 3.1: Centered reweighting equals the target policy return up to a constant
  • proof
  • Theorem 3.2: Chi-square penalty equals a $\chi^2$-divergence
  • proof
  • Proposition 3.3: On-policy (or near-on-policy) data implies no gain from correction
  • proof