Technical Summary: Zero-Shot Off-Policy Learning (ZOL)

1. Problem Statement

The paper addresses the challenge of Zero-Shot Off-Policy Learning within the paradigm of Behavioral Foundation Models (BFMs). BFMs are agents pretrained on large, unlabeled, reward-free datasets to learn reusable behavioral priors, allowing them to adapt to new tasks at test time without additional training (e.g., via text prompts, goals, or motion capture).

However, this inference-based adaptation inherits a critical failure mode from off-policy and offline Reinforcement Learning (RL): distributional shift.

• The Mismatch: A policy induced by a new task specification often queries state-action regions that are poorly represented or entirely absent in the pretraining dataset (Out-of-Distribution or OOD).
• The Consequence: Standard zero-shot inference (e.g., selecting a latent vector $z$ based on reward projection) leads to systematic errors in value estimation and value function overestimation. The agent may select latents that appear optimal under extrapolated values but induce visitation distributions unsupported by the data, resulting in suboptimal or unsafe behavior.
• The Gap: Existing solutions like DICE (DualDICE, OptiDICE) require task-specific optimization, which contradicts the "zero-shot" requirement. Other fast-adaptation methods (ReLA, LoLA) rely on online environment interaction and on-policy rollouts, which are costly and violate the fully offline nature of the BFM setting.

The core question posed is: Can we estimate optimal stationary occupancy ratios for a whole family of BFM-induced policies without per-task online optimization, enabling fast test-time adaptation while avoiding off-policy pathologies?

2. Methodology: Zero-Shot Off-Policy Learning (ZOL)

The authors propose ZOL, a training-free, test-time procedure that bridges the gap between Successor Measures (SM) and Stationary Density Ratios.

Theoretical Insight

The method relies on a theoretical connection derived from the Forward-Backward (FB) framework (Touati & Ollivier, 2021), which is used to estimate Successor Measures.

1. Successor Measures (SM): The SM, $M^\pi$, characterizes the discounted future visitation of a policy $\pi$. In the FB framework, the successor density $m_\pi$ is approximated via a low-rank factorization: $m_\pi(s_0, a_0, s, a) \approx F(s_0, a_0, z)^\top B(s, a)$, where $F$ is a forward embedding and $B$ is a backward embedding.
2. Density Ratio Identity: The authors derive an exact identity (Eq. 9) showing that the stationary density ratio $w_{\pi/\beta} = d_\pi / d_\beta$ (the correction factor needed to align the policy's visitation with the dataset) is proportional to the expected successor density under the initial distribution $\rho_0$:
   $$w_{\pi/\beta}(s, a) = (1-\gamma) \mathbb{E}_{(s_0, a_0) \sim \rho_0} [m_\pi(s_0, a_0, s, a)]$$
3. Low-Rank Estimator: By substituting the FB approximation into the identity, the density ratio becomes a simple linear function of the backward embedding:
   $$w_{\pi_z/\beta}(s, a) \approx (1-\gamma) W_{\pi_z}^\top B(s, a)$$
   Here, $W_{\pi_z}$ is the expected forward embedding over the initial distribution. Crucially, this estimator is data-only; it is evaluated solely on samples $(s, a)$ present in the dataset $D$, ensuring the policy never assigns probability to unsupported regions.

Algorithm

ZOL operates in three steps without modifying the pretraining of the BFM:

1. Pretraining: Train the FB embeddings ($F$ and $B$) using standard temporal-difference objectives on the offline dataset.
2. Latent Initialization: For a new task with reward $r_T$, infer an initial task embedding $c_T$ by projecting the reward onto the backward features: $c_T = \mathbb{E}_{(s,a) \sim D}[B(s, a) r_T(s, a)]$.
3. Latent Optimization: Optimize the latent vector $z$ to maximize a regularized weighted return:
   $$J_{total}(z) = \mathbb{E}_{(s,a,r) \sim d_\beta} [w_{\pi_z/\beta}(s, a) (r(s, a) - \bar{r})] - \lambda_1 \chi^2(z) - \lambda_2 L_{trust}(z)$$
   • Reweighting: The term $w_{\pi_z/\beta}$ reweights the supervision toward dataset-supported regions, correcting for the distribution shift.
   • Regularization:
     • $\chi^2$ penalty: Prevents extreme variance in importance weights.
     • Trust region ($L_{trust}$): Keeps the optimized latent close to the initial task embedding to prevent catastrophic drift.
     • Reward centering: Prevents trivial solutions (e.g., "do nothing" policies) caused by constant reward offsets.

The optimization is performed via gradient ascent (Adam) on $z$ using Monte Carlo sampling of start states to estimate $W_{\pi_z}$, requiring no environment interaction beyond the initial state sampling.

3. Key Contributions

1. Theoretical Bridge: The paper establishes a novel theoretical link showing that Forward-Backward successor representations implicitly contain the stationary density ratios required for off-policy correction. This allows for the computation of optimal correction weights directly from pretrained components.
2. ZOL Algorithm: A simple, training-free, and interaction-free test-time adaptation method that mitigates off-policy overestimation and extrapolation errors by explicitly enforcing support constraints via density ratio reweighting.
3. Empirical Validation: ZOL is evaluated across diverse benchmarks (ExORL, OGBench, SMPL Humanoid) and consistently outperforms state-of-the-art zero-shot baselines (FB, HILP, TD-JEPA) and online adaptation methods (LoLA, ReLA) in terms of robustness and return, particularly in long-horizon and high-dimensional tasks.
4. Ablation and Analysis: The paper provides a minimal 2D toy example demonstrating that standard zero-shot inference fails due to support bias, while ZOL successfully recovers task structure by aligning inference with the induced occupancy.

4. Experimental Results

The authors evaluated ZOL on three benchmarks:

• ExORL (DeepMind Control Suite): ZOL achieved consistent gains over baselines across Walker, Cheetah, Quadruped, and PointMass domains. Notably, it showed significant improvements in long-horizon tasks and high-dimensional domains (e.g., Quadruped), outperforming online methods like LoLA and ReLA without requiring environment rollouts.
• OGBench (Long-Horizon Goal-Conditioned RL): On tasks requiring temporal stitching (e.g., AntMaze, Cube manipulation), ZOL achieved superior success rates compared to FB and other baselines. It matched or exceeded the performance of online methods (LoLA/ReLA) while operating entirely offline.
• SMPL Humanoid (Motion Tracking): Using a pretrained FB-CPR model, ZOL improved motion tracking performance on the AMASS dataset, finding substantially better policies than the baseline FB-CPR inference.

In all benchmarks, ZOL demonstrated robustness to hyperparameter choices, requiring minimal tuning.

5. Significance and Claims

The paper claims to bridge the fields of off-policy learning and zero-shot adaptation. Its primary significance lies in demonstrating that robust zero-shot adaptation does not necessarily require expensive online interaction or per-task optimization.

• Mitigating the Inference Gap: The authors argue that performance gaps in BFMs often arise not from the representational capacity of the model, but from the suboptimality of the inference procedure itself. ZOL addresses this by ensuring the adaptation process remains grounded in the pretraining distribution.
• Training-Free Adaptation: Unlike methods that require fine-tuning or online rollouts, ZOL offers a "plug-and-play" solution that leverages the existing FB structure to correct distributional shifts at test time.
• Theoretical Justification: The work provides a principled justification for why FB representations are suitable for off-policy correction, moving beyond heuristic adaptations.

Limitations: The authors acknowledge that ZOL inherits the representational bias of the learned FB factorization. If the density ratio estimate is inaccurate or if a downstream task strictly requires sustained out-of-support behavior, the method's conservative nature (enforcing support constraints) may trade off performance for safety. Future directions include uncertainty-aware corrections and extensions to partially observable domains.