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.
