Zero-Shot Adaptation of Behavioral Foundation Models to Unseen Dynamics

TL;DR

The paper tackles zero-shot adaptation in Behavioral Foundation Models by addressing a key limitation of Forward-Backward representations: interference across unseen dynamics due to dynamics-agnostic successor measures. It introduces Belief-FB (BFB), a transformer-based belief estimator that conditions FB on an inferred context $h$, and Rotation-FB (RFB), which further partitions latent policy directions via a context-aligned von Mises–Fisher prior to reduce interference; this yields a bound on the worst-case error that is independent of the number of training dynamics. Empirically, BFB and RotFB outperform baselines on both seen and unseen dynamics across discrete and continuous CMDPs, with ablations showing the importance of context inference, trajectory length, and dataset diversity. The work advances practical zero-shot RL in settings with changing dynamics and partial observability, with potential impact on robotics and other real-world adaptive control tasks.

Abstract

Behavioral Foundation Models (BFMs) proved successful in producing policies for arbitrary tasks in a zero-shot manner, requiring no test-time training or task-specific fine-tuning. Among the most promising BFMs are the ones that estimate the successor measure learned in an unsupervised way from task-agnostic offline data. However, these methods fail to react to changes in the dynamics, making them inefficient under partial observability or when the transition function changes. This hinders the applicability of BFMs in a real-world setting, e.g., in robotics, where the dynamics can unexpectedly change at test time. In this work, we demonstrate that Forward-Backward (FB) representation, one of the methods from the BFM family, cannot distinguish between distinct dynamics, leading to an interference among the latent directions, which parametrize different policies. To address this, we propose a FB model with a transformer-based belief estimator, which greatly facilitates zero-shot adaptation. We also show that partitioning the policy encoding space into dynamics-specific clusters, aligned with the context-embedding directions, yields additional gain in performance. These traits allow our method to respond to the dynamics observed during training and to generalize to unseen ones. Empirically, in the changing dynamics setting, our approach achieves up to a 2x higher zero-shot returns compared to the baselines for both discrete and continuous tasks.

Figures (12)

• Figure 1: Summary of results. Aggregate mean performance over seen (train) and unseen (test) dynamics for zero-shot RL. The error bars indicate standard deviation over three seeds. Notably, both BFB and RFB adapt not only to the dynamics seen during training but are also able to generalize to unseen dynamics. There are 30 (20) training (test) dynamics for FourRooms and PointMass and 16 (4) for AntWind environments.
• Figure 2: Randomized-Doors environment for three different layouts, each produced through varying the grid structure (exact randomization procedure is a hidden variable) (left-middle) From state $s$, the goal of an agent is to capture a diamond at target location by picking up the most probable policy $\pi_z$ (yellow for the first type and purple for the second) to move to the closest open door based on internal representation. (middle) When there are multiple possible future outcomes in the training data from the same state, the $\pi_z$'s (different colors) interfere with each other, leading to picking up an averaged policy.
• Figure 3: Three different environment configurations from \ref{['fig:maze_robot']} are visualized (yellow, purple and mixed trajectories). For a fixed state $s$ and same goal across configurations, arrows depict latent directions $z_{\text{FB}} \in \mathcal{Z}$ and colored by optimal action as $a_{color}=\arg \max _a F(s, a, z_{\text{FB}})^T z_{\text{FB}}$. (left-middle) When FB is trained on the two distinct configurations in separation, most of the latent directions agree on the optimal policy $\pi_z$. (right) When FB is trained on mix of CMDPs and at test time tasked with any particular configuration from train, obtained policy is ambiguous, since most policy-encoding directions do not agree on the action.
• Figure 4: Visualization of inferred contexts $h$ from space of all possible contexts $\mathcal{H}$ (depicted as arrows) and task vectors $z_{\text{FB}}$ (depicted as points on sphere boundary). Transitions from same CMDP colored the same. Concentration parameter $\kappa$ defines spread of clusters. (left) Untrained transformer $f_{\text{dyn}}$ output for different transitions is unstructured and same transitions coming from same CMDP (identical colors) are not collinear. (middle) New sampling procedure aligns policy specific vectors $z_{\text{FB}}$ with context specific $h$, but clusters overlap before training. (right) After training, $h$ for transitions from the same context are aligned and policies $z _\text{FB}$ do not interfere between different environment configurations.
• Figure 5: Ablations on data diversity and context length of transformer encoder. We show the influence of number of environments (data diversity) and context length on train and test performance in Four-Rooms and Pointmass environments. For data-diversity ablation, we see a clear performance boost up until some point, after which it platoes, as the \ref{['eq:regret']} predicts. In our context‐length ablation, we observe similar behaviour: performance improves as the context grows up to the length of a single episode, and then levels off. The results are averaged across three seeds, the opaque fill indicates standard deviation.

Theorems & Definitions (7)

• Theorem 1: Regret-bound for Multiple Dynamics
• Theorem 2: Regret bound under latent space partitioning
• Theorem : Regret-bound for Multiple Dynamics
• Lemma 1
• proof
• Theorem : Regret-bound under latent space partitioning
• proof