Zero-Shot Adaptation of Behavioral Foundation Models to Unseen Dynamics
Authors: Maksim Bobrin
Ilya Zisman
Alexander Nikulin
Vladislav Kurenkov
Dmitry V. Dylov
Venue: ICLR 2026 Poster (2026)
Keywords: zero-shot reinforcement learning
unsupervised reinforcement learning
successor measure
Primary area: reinforcement learning
TL;DR: We provide both theoretical and empirical evidence that Forward–Backward representations cannot adapt to changing dynamics and introduce a method that overcomes this, generalizing to both seen and unseen dynamics at test time.

Abstract: Behavioral Foundation Models (BFMs) proved successful in producing near-optimal policies for arbitrary tasks in a zero-shot manner, requiring no test-time retraining 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 produce reasonable policies under distinct dynamics, leading to an interference among the latent policy representations. To address this, we propose an FB model with a transformer-based belief estimator, which greatly facilitates zero-shot adaptation. Additionally, we show that partitioning the policy encoding space into dynamics-specific clusters, aligned with the context-embedding directions, yields additional gain in performance. Those traits allow our method to respond to the dynamics mismatches 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.

Published as a conference paper at ICLR 2026
ZERO-SHOT ADAPTATION OF BEHAVIORAL FOUNDA- TION MODELS TO UNSEEN DYNAMICS
Maksim Bobrin
Ilya Zisman Alexander Nikulin
Vladislav Kurenkoy 2,1 Dmitry Dylov
'Computational Imaging Lab 2AXXX
MSS
*Innopolis University
HHumanoid
ABSTRACT
Behavioral Foundation Models (BFMs) have proven successful at producing near- optimal policies for arbitrary tasks a zero-shot manner, requiring no test-time retraining task-specific fine-tuning. Among the most promising BFMs are those that estimate the successor measure, learned in an unsupervised way from agsk-agoostic offline data. However, these methods fail to react to changes the dynamics, making them ineffective under partial observability or when the transition function changes. This hinders the applicability of BFMs in real-world settings (e.g., robotics), where the dynamics can change unexpectedly at test time. In this work, we demonstrate that the Forward- -aaaad-- (FB) representation, cannot produce reasonable policies under distinct dynamics, leading to interference among latent policy representations. To address this, we propose an FB model with transformer-batt based belief estimator, which greatly facilitates zero-shot adap- tation. Additionally, we show that partitioning the policy- colicy-encoding space into dynamics-specific clusters aligned with context-embeddingg directions yields fur- ther performance gains. These traits allow our method to respond to the dynamics mismatches observed during training and generalize to unseen dynamics changes. Empirically, in settings with changing dynamics, our approach achieves up to 2x higher zero-shot returns than baselines on both discrete and continuous tasks.
Project page: https: //tpl///lleeeeeee ief- fb....bb... dunnol 1 ab.ai/
INTRODUCTION
One very desirable property of reinforcement learning (RL) agents is their ability to adapt during inference to new tasks without requiring any additional learning. Achieving this in few trials as possible would be even better: the ideal being the zero-shot adaptation (Barreto et al., 2017; Touati al., 2022), where the agent never interacts with the environment at test-time and relies solely the task-agnostic data. Behavioral Foundational Models (BFMs) (Pirotta et al., 2023; Sikchi et al., 2024; Tirinzoni al.) may be considered step in this direction, because they can learn variety of behaviors from offline data without task specific labels. During inference, it is possible to extract task-specific policy that is theoretically optimal in terms of performance. Recent works (Tirinzoni et al.) demonstrates that methods based on successor measure estimation through Barward-Backwee (FB) decomposition (Touati & Ollivier, 2021), especially versatile and can successfully extract diverse policies from data.
At the same time. FB has fundamental drawback that limits its adaptation ability to more general settings. In our paper, we show that FB is unable to generalize across different environment configu- rations (i.e., dynamies), such as changes in a transition function (e.g., new obstacles or environment configuration) some latent factor variation (e.g., wind direction change in mass). This limitation stems from the way the approximation of the successor measure (Blier et al., 2021) is estimated: FB averages the discounted aiture-occupancyy state distribution over all observed dynamics, which inevitably causes interference in a policy representation space. This fact alone may severely constrain the applicability of FB in the real-world scenarios. For example, one of the largest robotics dataset, Open X-Embodiment (Collaboration, 2023), consists of 22 different robot embodiments, and training
corresponding author: bobr 8 aomail-0 semail.com work is done collaboration with dunnolab.ai

---

Published as a conference paper at ICLR 2026
FourRooms
PointMass
AntWind
1.0
740
1.0
0.76
750
680
0.7
Train
0.5
0.4
0.5
0.45
0.2
500
410
340
390
0.2
0.25
0.15
250
0.0
0.0
0
0.6
0.61
0.6
0.55
0.45
600
640
550
Test
0.4
0.4
0.4
400
410
290
250
0.2
0.2
0.1
0.15
0.2
0.1
0.1
200
0.0
0.0
0
RFB (ours)
BFB (ours)
HILP
LAP
FB
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 5 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 Cettomiize--ourRRooooo and PointMass and 16(4) for Ant Wind environments.
FB on each them independently is infeasible. In Section 3.1, we discuss the underlying cause in more details by providing both empirical and theoretical support.
To remedy such limitation, we introduce Belief-FB (BFB), a conditioning method for FB through belief estimation, popular technique of uncertainty quantification in Meta-RL (Zintgraf et al., 2020; Dorfman et al., 2021). To implement this, we employ a permutation-invariat transformer encoder, denoted Saya. which processes given trajectory from the dataset produce dynamics- encoding vector This representation subsequently utilized as a conditioning input to the successor feature network, expressed ( h, -). We pre-train Jáyn self-supervised fashion, thus posing no additional requirements on the data structure or the trajectory re-labeling, while maintaining theoretical guarantees. We discuss the implementation of Belief-FB Section 3.2.
Remarkably, Belief-FB enables the generalization capabilities of FB not only through the dynamics seen in the training dataset, but also on the unseen test configurations. Additionally, we find that in order to align belief estimation better with FB, one also needs to partition the policy space encodings prior into dynamics-specine clusters, which significantly improves generalization abilities of FB. We propose Rotation-FB (RFB) that accomplishes this partitioning. We present the theoretical support and the implementation details of Rotation-FB Section 3.3. Empirically, both BFB and RFB outperform baselines for seen and unseen dynamics, as gathered in Figure and discussed in Section 4.3.
We believe that our work sufficiently broadens the possible applicability of BFMs, yet keeping all of the zero-shot properties unchanged. Our contributions are as follows:
We demonstrate the limitation of Berward-Backward (FB) representations (Touati Ollivier, 2021), which lies in its inability to generalize per across different dynamics both from train and test, where dynamics shift constitute new layout grids or changes in the transition function that are hidden from an agent. Refer to Section 3.1 for more discussion.
We propose Belief-FB (BFB), which employs permutation-invarian transformer encoder to infer belief over the current dynamics (Zintgraf et al., 2020; Dorfman et al., 202 1). Analyzing BFB's policy encoding space reveals that additional disentanglement is benefi- cial, motivating our Rotation-FF (RFB) extension. Section 3.2 examines Belief-FB, and Section 3.3 details Rotation-FBBs theoretical motivation and implementation.
2

---

Published as a conference paper at ICLR 2026
We empirically demonstrate that both BFB and RFB can adapt to different dynamics, unlike its counterparts in the zero-shot setup. Refer to Section 4.3 for the discussion and Figure 1 for results.
2 BEHAVIORAL FOUNDATION MODELS (BFMs)
A Behavioral Foundation Model (BFM) (Pirotta et al., 2023; Tirinzoni et al.; Frans et al., 2024; Park al., 2024; Sikchi al., 2025) is an RL agent trained in an unsupervised manner on a task-agnostic dataset to approximate optimal policies for diverse reward functions (tasks) specified at inference (test-time).
Forward- Bacddddd Representation (FB) (Touati & Ollivier, 2021) approximates discounted succes- sor measure (Blier et al., 2021) for various behaviors across diverse tasks. The successor measure M (So. 0o, X) for subset X C S is defined as cumulative discounted time spend at X starting at (so, ao) and following π thereafter, defined as (for finite state-action space case)
M 00, X)
Σ
7 PP((0+ E X ||00, ao T
120
with the corresponding Q-function for a specific task r:
O- (s0, ao) = Σ
st)
(2)
--x
In the continuous case, the FB representation aims to approximate successor measure through finite- rank approximation under diverse policies through forward F S x A 2 - R and backward B B:S + Rd functions. Given a set of policies T= parametrized by task variable drawn uniformly from sphere 2FB E Unif(2 Z = St). Assuming P is probability distribution over states within the offline dataset, the objective for FB is written as M " (so, ao X) 2 Sstex F(50, 00, z)" B(s )((ss). Then, policy can be extracted as
(s) arg max F(s, a, =)
For continuous case, the greedy policy is approximated via DDPG (Lillicrap al., 2015). We refer to the Appendix B.1 for in-depth details regarding FB training procedure. During test time the task policy parametrization is approximated as Stest a)spa))p (s, B(, a)]. If the inferred task vector Stest lies within the task sampling distribution (in linear span) of Z used during training, then the optimal policy for task Ttest is obtained from Equation 2 as (s(s) as arg max, Q (8, a). Extended discussion on other related works is included in the Appendix A.
3 METHOD
Problem Statement. The dataset consisting of diverse environments can be formally considered as Contextual Markov Decision Process (CMDP) defined by a context space C and a mapping M e E C C M(c) (S, A, re: 7), where both S, A are shared across contexts, Te S x A + (SS) is the context-dependent transition kernel, r. is the reward function, Pa E (SS) is the initial state distribution, and 7 [0, 1) is the discount factor. Each context (e.g., wind direction, friction, or door locations) specifies a unique MDP.
When the context is unobserved, the problem becomes a POMDP. Under standard assumptions, there exists a sufficient distory-dependerd statistic-the belief state bx(c) = PP He) E ACC)-capturing the posterior over contexts given the history He. Solving the POMDP is equivalent to solving the fully observable Belief-MDD (A(C) x S.A, Pb. 7), where states are augmented with beliefs.
We assume access to an offline, reward-free dataset Dussin consisting of trajectories {(Sk: ag2 1 collected under diverse exploratory policies from a finite set of training contexts Cirain C. At test time, for an unseen context Gea E C Cirans we are given a short reward-
history H- = {(st, at ( from an exploratory policy in M(Cres), and task specified by
reward r : A
The goal is to infer an approximate belief b(e H) for any given history (trajectory) and extract zero-shot policy и (without additional learning) that minimizes the regret
R
sup EC Cinis
E (ssa))paan
Q (s, a) Q7 (s,a).
(4)
We use the "zero-shot RI." following Touati & Ollivier (2021).
3

---

Published as a conference paper at ICLR 2026
Figure 2: Sandomized-Doorr environment for three different layouts, each produced through varying the grid structure (exact randomization procedure is a hidden variable) (eff-middle) From state s, the goal of an agent is to capture a diamond at target location by picking up the most suitable policy яı (yellow for the first type and purple for the second) to move to the closest open door based on internal representation. (right) When there are multiple possible future outcomes in the training data from the same state, the T 's (different colors) interfere with each other, leading to picking up averaged policy.
where n²" is the optimal policy for task r under dynamics Ta..: To formally study optimality guarantees of the problem above, we employ the following assumption commonly used for dynamics generalization (Eysenbach et al., 2021; Jeen & Cullen, 2024):
Assumption 1 (Coverage). The test initial state-action distribution Pres is supported on the support of p, i.e. supp(pre) S supp(p) (equivalently, Prest < p. i.e absolutely continuous).
3.1 INVESTIGATING LATENT DIRECTIONS SPACE UNDER MULTIPLE DYNAMICS
We begin by addressing the following question: Why does FB representations fail generalize effectively to different situations under dynamics variations, i.e., if learned on data sampled from diverse CMDPs? While the answer may appear intuitive, a closer look into the geometric structure of learned latent directions 2FB E Z, which encode possible policies π. reveals critical insights which will be helpful later. We approach this question both theoretically and empirically on custom didactic discrete oarteally-bbservable Randomized Doors (see Appendix C.1) environment for building intuition. Partial observability adds additional challenges and showcases the need estimate belief state, which we discuss in the following sections.
In this experiment the only source of dynamics variation is the grid layout type. Namely, the positions of doors and walls are changed each new episode, depending hidden configuration variable c. We collect dataset of random trajectories drawn from multiple layouts, yielding near-uniform coverage of the entire state-action space. Now, consider particular state that agent finds itself in three different layouts (see Figure 2). During FB training, we estimate expected successor features via F(s,', , 2rB) for policy representations 2FB Uniform(Sd-) starting at s.
In this setting, optimal successor features require different optimal policies, depending on the layout an agent is instantiated in. Because 2h8 does not enforce separation (i.e., prior) over layout- specific futures, the FB model suffers from interference: latent directions encoding conflicting future outcomes overlap and become entangled in the policy representation space Z. For each layout configuration and fixed state 8, Figure 3 depicts latent directions žpg, colored by optimal policy as a,,,,,/))) 2FB)" zpB. When FB is trained on first two layouts separately, a unique dominant behavior (colored) emerges in Z, recovering the optimal goal-reaching policy π. This means that any randomly sampled 2 at training time agrees on the possible future (successor features). In contrast, training on mixed data, where transitions come from multiple environment instances, causes SFB to blend dynamics-specifte information and average over possible futures, yielding a policy that is suboptimal for every layout, including those in the training set. These observations are theoretically supported by the following:
acolor arg max, F(s, a,
Theorem 1 (Regret bound via uniform successor approximation under Assumption 1). Let r be bounded with ||||| < R with R sup( (r(s,a) and discount 7 E (0, 1). For any test
4

---

Published as a conference paper at ICLR 2026
Z1
Z2
2mix
Figure 3: Visualization of the randomly sampled vectors FB, which encode policies Tam for three datasets from Figure 2. For a fixed state 8 and same goal state across config- urations, arrows depict latent directions FFB
and colored by dominant directions as Acolor arg max, F(s, a, =FB)" 2FB. (left-middle) When FB is trained on the two configurations independently, most of the latent directions agree on the optimal policy म.. (right) When FB trained on mix of CMDPs and at test time tasked with any particular configuration from train, obtained policy ambiguous, since most eacccy-nncoding directions do not agree on the action.
CMDP satisfying the coverage assumption, the policy næ returned by the method obeys
((s,))~Paa (s, (s,aa) s 1 1-7
R (Ek Dest).
We provide full proof in Appendix B. Because ck+1 <g by monotonicity of the worst-case approximation error over a fixed function class, the upper bound equation becomes looser as included training time. This statement concerns the approximation term
more environments only. In practice, adding CMDPs may also increase the dataset size and reduce the finite-sample estimation term (see additional discussion in the Appendix B), so the net effect on regret empirical.
Futher, in Section 3.3, we refine this result and show that the explicit dependence on the total number of environments thereby tightening the upper bound when kmax < k.
be replaced by dependence on kmax (the size of the largest cone/cluster),
This interference highlights fundamental trade-off. FB expressive enough to model any task in the linear span of the reward, and yet, when trained across environments with distinct unobserved parameters, the lack of contextual conditioning forces it to average expected future across different dynamics rather than separate them. The resulting successor measure merges transitions from distinct layouts and entangles directions in the latent space 2. To disentangle these directions, we must represent uncertainty about the hidden context explicitly. The next section introduces belief- conditioned objective that infers the latent context and allows FB to maintain environment-specifie successor measures.
Takeaway
Because FB training inherently averages over all possible successor features, it cannot learn disentangled policy space and. therefore, fails adapt changes in dynamics.
3.2 BELIEF STATE MODELING
To resolve the interference issue described in Section 3.1, we infer the latent context of envi- ronment and augment FB input that belief. We train transformer encoder Sayn. by passing to set of transitions {(tt, at, 1) and outputting an h E Rdd We denote the space of all possible inferred contexts as H. where each element h encodes dynamics for particular environment. Because the ordering is discarded and no rewards in transitions provided, the encoder must focus on dynamics specific mismatches (e.g., layout geometry, friction or wind direction), rather than policy specifies. Such context encoder should be permutation invariant, since unobservable factors describing environment are independent of the order of transitions in an episode. This setting provides theoretical ground for zero-shot and fho- learning Snell al. (2017).
5

---

Published as a conference paper at ICLR 2026
H
2 Rotation
Figure 4: Inferred context vectors and task vectors across training stages. Arrows represent inferred contexts h E H: points on the sphere boundary represent task vectors 2Fh (middle panel). Points are colored by CMDP, so transitions from the same CMDP share the same color. The concentration parameter x controls the cluster dispersion. (leff) Untrained Sayn: inferred contexts are unstructured; transitions from the same CMDP not aligned. (middle) New sampling procedure (before training): 2FB is encouraged to align with the corresponding h, but clusters across CMDPs still overlap. (right) After training: contexts from the same CMDP form coherent, aligned clusters, and task vectors 2PB separate across environment confi configurations reducing interference.
Concretely, dataset consists of episodes ([(*). Gz 1)4)), } coming from CMDPs with randomly instantiated hidden specification variable c (different dynamics). We train a transformer encoder on random episodes (without episodic labels q)) of context length n to infer contextual (hidden) variable which fully specifies the dynamics across given episode. The transformer encoder loss involves two main components: 1) h is encouraged to follow Gaussian prior and is shared across trajectory. and 2) projection head, which combines h with (s (, a.) to predict 1 Those stages can be either trained end-to-end or separately. We observed that separating FB training from Sayn gives better results.
For each trajectory we concatenate the inferred context vector h with the task vector 2FB to obtain augmented input [h; 2FB] and condition only forward network as:
Mr, (st, dt $1+1) = (st, 4r, [h; zrB)) " B(s+1).
(6)
We empirically found that conditioning the backward network B degraded performance, producing smoothed out Q function, so remains shared across contexts. Algorithm is summarized in Algorithm 1.
At test time, the agent is provided with a short (context length), reward-free trajectory and is passed Sayn to obtain h. By plugging the result into Equation 3, the greedy policy obtained.
Takeaway 2
We train a transformer in a sulf-super regime to estimate a belief over possible contexts, augmenting FB inputs and enabling effective disentanglement of contextual representations.
3.3
STRUCTURING DIRECTIONS IN THE LATENT SPACE
Insights from Section 3.1 showed that sampling task-vectors 2FB uniformly on the hypersphere encodes averaged policies, while Section 3.3 provided a solution through explicit context identification. We now combine these observations together through enhanced sampling žpg around the inferred context h.
In Vanilla-FB, each state s draws 2FB ~ Uni with no inductive bias, so resulting policies n, conflict with each other in CMDP setting, even if additional explicit conditioning is introduced as before. We replace uniform prior with a von Mises-Fisher (yMF) distribution centered at the context direction for episode h = faya (((si, ai 8141)))
(7) with x controlling the spread diversity of policies (left and middle figures from Figure 4). In practice, to draw rn we first pick a simple vector (e.g.. the first basis vector), perturb with vMF noise, and finally rotate the result onto h with Householder reflection.
24+rB ~ vMF(u = h, k).

---

Published as a conference paper at ICLR 2026
This enhancement has several benefits: 1) because directions h that differ in dynamics occupy disjoint cones on the hypersphere, FB can fit the successor measure locally inside each cone, avoiding the destructive averaging effect quantified in Section 3.1 and 2) alignment procedure encourages the agent to explore policies that are plausible under its current belief while still injecting controlled diversity through
Importantly, such procedure also lowers the Theorem upper bound by replacing its dependence on the total number of environments k with a dependence on kmax (the size of the largest cone).
Theorem 2 (Regret bound under latent-space partitioning). Let hi,... N E s be the context directions from Sdyn and let {Cz}á=ı be disjoint cones around them. Assume block-separable parame- terization (Assumption in Appendix B), so that losses from & E C2 depend only on block (Fj, By).
kmax max; Ki C. then
6 =
max Sic,1 S Ekmaxe 1SISL
and Theorem 3 holds with Et replaced by max, Fic,r (See Appendix Theorem 4.)
Intuitively, Theorem 2 states that after the partitioning procedure of the latent space into non- overlapping clusters based on context representations f, the global worst-case FB approximation error é& = maxjgr Ej is determined only by the cluster whose error éj is largest. Importantly, the bound depends on kmax rather than the total k. When kmax is controlled (e.g., via non-overlapping cones induced by an appropriate concentration k), the bound becomes effectively independent of k. Full proof can be found in Appendix B
Takeaway 3
Adjusting the prior over task vectors 2FB further mitigates the averaging effect and disentan- gles policy representations better based on the inferred dynamics.
4 EXPERIMENTS
In this section, compare proposed methods, namely: Belief-FB (BFB) (Section 3.2) and its extension Rotation-FB (RFB) (Section 3.3), against the baselines in discrete and continuous settings. We outline experiments design below; all other necessary details are provided in Appendix D. Every environment is framed as a contextual MDP (CMDP), where the context differs by the underlying hidden variation (e.g., grid layout, transition dynamics). During test time, we provide single trajectory from random exploration policy, which enables context inference.
4.1 ENVIRONMENTS AND SETUP
To support claims and theoretical insights made in previous sections, we consider the following experimental setups: (i) discrete, partially observable Randomized Four-Rooms (Appendix C.2), (ii) continuous AntWind (Appendix C.3), and lastly (iii) continuous partially observable Randomized- Pointmass (Appendix C.4). We vary the number of train layouts for each experiment, while fixing the number of held-out unseen context settings to 20 for Randomized Four-Rooms and Randomized- Pointmass, and for Ant-Wind. We perform comparisons against following baselines:
HILP (Park et al., 2024) is a method that learns state representations from offline data so that the distance in the learned representation space is proportional the number of steps between two states in original space. FB (Touati & Ollivier, 2021) is an original version of the FB, described in Section 2. Laplacian RL (LAP) (Wu al., 2019) constructs graph Laplacian over state transitions from experience replay, then computes its eigenvectors to form low-dimensional representations that capture the environment's intrinsic structure. Random agent, which randomly explores the environment in mnsk-independent manner.
Randomized Four-Rooms is a discrete, deterministic, partially observable environment, where the task is to optimally move to the goal location. Training data is collected by executing random policies in N distinct grid layouts, that differ in doorway and wall locations.
Ant- Wind is continuous environment, where the goal is to make an ant to walk forward as fast as possible. The environment dynamics are determined by the direction (angle) of a wind d.
7

---

Published as a conference paper at ICLR 2026
Randomized Four-Rooms
Pointmass
Train Return
1.0
RFB
0.75
BFB
HILP
0.5
Lap
FB
0.25
Random
0.0
Test Return
1.0
0.75
0.5
0.25
0.0
10
20
30
50
50
100
150
200
10
20
50
50
100
150
200
# Environments
Context Length
of Environments
Context Length
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 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, the Theorem 1 predicts. In 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.
Randomized Pour-Rooms
Cariamizzed-Pintmasse is partially observable continuous environment, where the task is to move to the goal locations. Maze grid structure is generated randomly, where each cell either contains wall or empty, while ensuring there is a path between start and goal locations.
1.0
Test
Train
Return
0.75
0.5
0.25
4.2 COULD
BELIEF ESTIMATION ENABLE ADAPTATION FB?
0.0
Pointmass
Previously, we provided the theoretical foundations and spec- ulated on the matter why FB is unable to differentiate between distinct dynamics and how we can use the belief estimation to overcome this. We refer to Table and Figure that show our empirical findings to support
1.0
Return
0.75
claims.
0.5
would like to point out that neither FB nor LAP are able to outperform simple random baseline in PointMass and Four- Room, indicating that the policy they learn is most likely stuck in some obstacle due to averaging (see Section 3.1. Only HILP, which uses a different way to learn policy representations, is able to perform better than random policy.
0.25
0.0
10
30
50
100
K
Figure 6: Influence of x in RFB on performance. The results are aver- aged across three seed, the opaque fill represents standard deviation.
Belief-FB and Rotation-FB outperform every baseline method, indicating that belief estimation is indeed missing piece for adaptation. Notably, our methods also demonstrate generaliza- tion capabilities beyond train data on unseen test tasks.
4.3 DO BFB AND RFB CAPTURE HIDDEN PROPERTIES OF THE ENVIRONMENT?
For an agent to refine its policy, needs to keep track and update the uncertainty over possible environment configurations. Both Belief-FB and Rotation-FB accomplish this. Figure 7 illustrates this phenomenon visually. In Randomized- Door (left), the episodic trajectories from five layouts form non-overla clusters in the first two principal components of h, effectively disentangling different dynamics.
In Ant-Wind. the embeddings lie almost perfectly circle whose azimuth matches the underlying wind direction, generalizing smoothly to the 4 held-out wind angles. The quantitative results for evaluation in Table 2 (averaged across all environments) reveal that the baseline methods fail to
8

---

Published as a conference paper at ICLR 2026
Randomized-Dowe
AntWind
Layout
Environment
t-SNE dimension
PC
Test
PC 2
[-SNE dimension
Figure 7: 2D projections of Edy inferred from different trajectories across number different contexts (colors), showing effective disentangling environments based on transition function or other mismatches. (left) First two principal components are visualized for estimated Zdyn from five trajectories, each representing different layout type in Randomized-De Dooors. (right) Inferred context variables Edyn recover hidden wind direction parameter AntWind environment both for train and test, proving successful extrapolation properties.
recover those environment-specific properties and therefore produce sub-optimal policies even for train cases. In particular, HILP tends to predict an average direction in Randomized Four-rooms and ignores obstacles, while FB outputs same policy and Q function for almost all environments. Figure 12 shows that Q function is properly estimated only for BFB and RFB, respecting wall positions.
4.4 DOES CHINNGE IN CONTEXT LENGTH INPUT TO fUYN IMPACTS PERFORMANCE?
In this experiment, we examine whether increasing the input trajectory length of improves perfor- mance. We vary the context length of fayn from 50 to 200 and present the results in Figure for both Randomized Four-Roomy and Randomized Pointmass environments, across train and test configu- rations. The results show that performance is poor when the context length is shorter than a single trajectory episode (100 steps), as short trajectories only capture local, near-term goals. Conversely, excessively long sequences provide no additional benefit due to redundancy, since fayn already contains all neccessary information. Evaluations both train and test environments demonstrate that Sayn produces representations capable distinguishing between different context instances while maintaining robustness.
4.5 DOES INCREASE IN DATASET DIVERSITY MAKE POLICIES MORE ROBUST?
We investigate if diversifying CMDP training configurations improves performance. Intuitively, broader state-action space coverage enhances successor measure estimation. Experiments confirm this: Figure 5 shows rapid improvement for BFB up to 25 configurations, while baselines match random policy performance. Once learned representations from sayn cover all variation modes (contexts), additional data yields minimal gain (< 3%). These results align with Theorem 1.
4.6
How K IN RFB INFLUENCES PERFORMANCE?
As described in Section 3.3, RFB concentration к regularizes the diversity of policies for each environment. One the one hand, concentration should be high ensure non-overlapping policy parametrized clusters  for different h, while at the same time it should not exceed certain value to control the diversity of policies in the environment, preventing collapsed solutions. Figure shows that lower values of k, meaning task-vectors likely producing overlapping clusters. As grows, task-vectors become more specialized, lowering variance which results in higher performance.
sampled with high deviation around h,
CONCLUSION & LIMITATIONS
In this paper, introduce Belief-FB (BFB) and Rotation-FB (RFB). two methods that extend Backward representations to handle dynamics mismatches, bridging the gap towards truly adaptive agents. At first, identify critical limitation in existing BFMs both theoretically
the Forward- Forward-Backward

---

Published as a conference paper at ICLR 2026
Table Zero-shot performance across environments with varying dynamics. Results for FourRooms, PointMass, and AntWind are aligned with the main paper. We add Oracle-ID (one-hot environment ID concatenation) and Contextual-FFB (our reimplementation of Jeen & Cullen (2024)). Oracle-ID excels in-distribution but fails to generalize out-of-distribution (OOD). Contextual-FB underperforms due to reliance on classifier expressivity. For the new OGBench Scene environment, we vary friction from 0.4-1.0 4.1 (train) and test on unseen low friction 0. 1-0.3, demonstrating dynamics generalization akin to Win (wind direction variation). Higher is better.
Method
Four Rooms
PointMass
AntWind
OGBench Scene
Train
Test
Train
Test
Train
Test
Train
Test
FB
0.05
0.15
0.04
0.20 +
0.10 0.03
390 +:
0.40 + 0.06
0.20 0.05
LAP
0.20 + 0.04
0.10
0.03
0.15 +: 0.04
0.10 + 0.03
340 1: 35
290 : 25
0.30 +: 0.05
0.10 1: 0.03
HILP
0.40 + 0.06
0.20
0.05
0.45 + 0.06
0.25 + 0.05
410 + 45
410 + 40
0.50 + 0.07
0.30 + 0.06
Contextual FB
0.05
0.18
0.04
0.15 0.04
0.60 ::
0.40 0.07
Oracle-ID
0.90 = 0.03
0.10 2: 0.03
0.92 + 0.02
0.08 = 0.02
780 +: 30
50 = 20
0.95 - 0.02
0.0 : 0.02
BEB (ours)
0.70 = 0.07
0.40
0.06
0.76 1: 0.07
0.45 + 0.06
680 1: 60
550 1: 50
0.6 =: 0.07
0.45 +: 0.06
RFB (ours)
0.04
0.61 0.05
0.05
0.55
and empirically: interference arises when training procedure relies on naively sampling policy- encoding latent directions for transition samples coming from conflicting dynamics. To address this, we learn hidden context variables (belief states) via a transformer encoder and use them as additional conditioning (Belief-FB) to the FB. Further, we improve latent- direction sampling by aligning task-relevant abstractions with environment-specife context, ensuring non-overlapping distinct environment specific regions in policy-encoding latent space. Both BFB and RFB demonstrate theoretical and empirical improvements over prior methods. However, limitations include evaluations on a narrow set of dynamics mismatches and the introduction of the additional hyperparameters, such x. that controls policy diversity across environments. Moreover, both BFB and RFB inherit all of the drawbacks of the vanilla FB representations: limitation to only linear reward representations and convergence guarantees. Also, random exploration at test time could fail at more complex environments and combining BFB and RFB together with more clever exploration methods at test time (Grillotti al., 2024; Urpí et al., 2025) would make methods more scalable.
As future research directions, it would be valuable to investigate whether other zero-shot RL methods, those not based on successor- meecese estimation, exhibit similar interference issues, and scale our approach to more complex benchmarks such as -Lnnd-MiniGrid (Nikulin et al., 2024; 2025) or Kinetix (Matthews al., 2025).
ACKNOWLEDGMENTS
This work was supported by the The Ministry of Economic Development of the Russian Federation in accordance with the subsidy agreement (agreement identifier 10000003113225P440000 grant No 139-15-2025-012..
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Siddhant Agarwal, Harshit Sikchi, Peter Stone, and Amy Zhang. Proto successor measure: Rep- resenting the space of all possible solutions of reinforeement learning, 2025. URL https: //apenreview. net/ tere/fooume forum? 9SV1iwocc..
Andre Barreto, Will Dabney, Remi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Has- selt, and David Silver. Successor features for transfer in reinforcement learning. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL https: / 2017/file 2 http://iee.Sbbbllllll12222233333 408s54444dde99beb05o Paper.
http:/:://cccceepppeeeeeee neur iss cc cc/paper_fillieeceeeeeeeeee /aper//
Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. A survey of meta-reinforcement learning, 2024. URL htttps://arxtiti org/ abs/
2301. 08028.
Steffen Bickel, Michael Brückner, and Tobias Scheffer. Discriminative learning for differing training and test distributions. In Proceedings of the 24th international conference on Machine learning. pp. 81-88, 2007.
Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal-dependent values: A mathematical viewpoint. arXiv preprint arXiv:2 2021.
Diana Borsa, Andre Barreto, John Quan, Daniel J. Mankowitz, Hado van Hasselt, Remi Munos, David Silver, and Tom Schaul. Universal successor features approximators. International Conference Learning Representations, 2019. URL https:: htttps://openreee/ //tp//ppeeeeeeeeeeeeeeeeeeeeeeee ew net/ forum? Ad=sivWjiRcKS
Open X-Embodiment Collaboration. Open X-Embodiment: Robotic learning datasets and RT-X models. https: /apps://arxxee /twrxiivvvoo.. org. a 2211 0.1.0.00. 08864, 2023.
Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural computation, 5(4):613-624 1993.
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta learning of exploration, 2021. URL https://azzi o als 2008 02598.
Benjamin Eysenbach, Shreyas Chaudhari, Swapnil Asawa, Sergey Levine, and Ruslan Salakhutdinov, Off-dynamics reinforcement learning: Training for transfer with domain classifiers. In International Conference on Learning Representations, 2021. URL https: //tttpss///peeevveee net forrmmmm?
Id=eqBWg3ACIAA, 3ACIAX.
Basura Fernando, Amaury Habrard, Mare Sebban, and Tinne Tuytelaars. Unsupervised visual domain adaptation using subspace alignment. In Proceedings of the IEEE international conference computer vision, pp. 2960-2967, 2013.
Kevin Frans, Seohong Park, Pieter Abbeel, and Sergey Levine. Unsupervised zero-shot reinforcement learning via functional reward encodings. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 13927-13942. PMLR, 21-27 Jul 2024. URL https: pptpe//dccccee......o..... mlr. wrwpsppsppppppppee / frans2 a html.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07221 2020.
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for rl. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bur Buce E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: http::/:://///cc/ecceeeeeeee/e/e neurips.. ce/ ac/paper_... httpp/perrfeeeee///eeeppaaaaeeee /tt.2/20444477231111277708880800 - pddf
files/ ppppe////p////////////

---

Published as a conference paper at ICLR 2026
Luca Grillotti, Maxence Faldor, Borja G León, and Antoine Cully. Quality-diversity actor-critic: learning high-performing and diverse behaviors via value and successor features critics. arXiv preprint arXhv:2403.099300 2024.
Tuomas Haarojja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pp. 1861-1870 Pmlr, 2018.
Scott Jeen and Jonathan Cullen. Dynamics generalisation with behaviour foundation models. In Workshop on Training Agents with Foundation Models RLC 2024, 2024. URL https://
openreview net/ fenttt/formmmmiddddd YM7vuP.
Scott Jeen, Tom Bewley, and Jonathan Cullen.
Zero- reinforcement learning from low quality data. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL wwwphhtp////o/o /ttp:::://ppeneeeee... net
net forum? id= (edd-799Wwkkiik
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research, 76: 201-264, 2023.
Wouter Kouw and Marco Loog. A review of domain adaptation without target labels. IEEE transactions on pattern analysis and machine intelligence, 43(3):766-775, 2019.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, Maxime Gazeau, Himanshu Sahni, Satinder Singh, and Volodymyr Mnih. In-context reinforcement learning with algorithm distillation, 2022. URL https::////omm /tattp:://aiiiiiii/ wwwaa.iiivv.... org a2w.
Jonathan N. Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning, 2023, URL
https: http///axvvvvve///e org/ abs / 2 14892.
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.022 02971, 2015.
Michael Matthews, Michael Beukman, Chris Lu, and Jakob Nicolaus Foerster. Kinetix: Investigating the training of general agents through open-ended physics-based control tasks. In The Thirteenth International Conference on Learning Representations, 2025, URL https: //wpr//vie......eeemmmmmm net/
Aditya Modi, Nan Jiang, Satinder Singh, and Ambuj Tewari. Markov decision processes with continuous side information. In Algorithmic learning theory, pp. 597-618. PMLR, 2018.
Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, Artem Agarkov, Viacheslav Sinii, and Sergey Kolesnikov. Xland :andd-miiggri:: Sealable meta-reinforrcement learning environments in jax, 2024. URL https:ttt///om / /tt///iivvv///////////////////av org/ abs/ 2212. 12044.
Alexander Nikulin, Ilya Zisman, Alexey Zemtsov, and Vladislav Kurenkov. Xland-1000 A large- scale multi-task dataset for in-context reinforcement learning, 2025. URL https: // /arxiv. orgg btg//bs/// 2406 6..00060000 08973.
Seohong Park, Tobias Kreiman, and Sergey Levine. Foundation policies with hilbert represen- tations. In Forty-first International Conference Machine Learning, 2024. URL https: //openrevies new. net. forum? d=L.NNsaAKuu.
Matteo Pirotta, Andrea Tirinzoni, Ahmed Touati, Alessandro Lazaric, and Yann Ollivier. Fast imitation via behavior foundation models. In NeurIPS 2023 Foundation Models for Decision Making Workshop, 2023. URI hnttpss://ppeeee /pw.eeeiee..eeeeeee/eee/ee.eeeee net. forum? d=SHNjk441 in.
Andrey Polubarov, Nikita Lyubaykin, Alexander Derevyagin. Ilya Zisman, Denis Tarasov, Alexander Nikulin, and Vladislav Kurenkov. Vintix: Action model via in-context reinforcement learning, 2025. URL htttts://aaaeeeeeeeeeeeeeeeeetet wwwp://wr./////oo/ /2501 1940.
12

---

Published as a conference paper at ICLR 2026
Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine Efficient off- policy meta-reinforreement learning via probabilistic context variables, 2019. URL https:
///axxiv.org/ab//115/3... 08225
Harshit Sikchi, Siddhant Agarwal, Pranaya Jajoo, Samyak Parajuli, Caleb Chuck, Max Rudolph, Peter Stone, Amy Zhang, and Scott Nickum. RI zero: Zero-shot language behaviors without any supervision. arXiv preprint urXiv:24/2.05718, 2024.
Harshit Sikchi, Andrea Tirinzoni, Ahmed Touati, Yingchen Xu, Anssi Kanervisto, Scott Nickum, Amy Zhang, Alessandro Lazaric, and Matteo Pirotta. Fast adaptation with behavioral foundation models. arXiv preprint 2504x:250000786 2025.
Viacheslav Sinii, Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, and Sergey Kolesnikov. In-context reinforcement learning for variable action spaces, 2024. URL /ttps://arxiv. orgg/bbs/2312..133
Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning, 2017. net forrmm? i -hh=-11Hhhllg
URL htkLhtttps///ppeeeee
Casper Kaae Sønderby, Jose Caballero, Lucas Theis, Wenzhe Shi, and Ferene Huszár. Amortised map inference for image super-resolution. arXiv preprint arXiv:1610.04990 2016.
Denis Tarasov, Alexander Nikulin, Ilya Zisman, Albina Klepach, Andrei Polubarov, Lyubaykin Nikita, Alexander Derevyagin, Igor Kiselev, and Vladislav Kurenkov. Yes, q-lcarning helps offline in-context RL. In Scaling Seproiig Foundation Models without Human Supervision, 2025.
URL httptttpp///.c /ttp:////ppereeeee http//peeevviiieeeeeeeeeeeeeeeee reweew.eeww...... net www.teettoouueeeeeeee forum? i d8 6JHZZUn.
Jayden Teoh, Pradcep Varakantham, and Peter Vamplew. On generalization across environments in multi-objective reinforcement learning. In The Thirteenth International Conference on Learning Representations, 2025. URL https: httpas///ppeeeiieeeeeeee net forum wattorumii=tuueee / 21d=tuE242UUQS
Andrea Tirinzoni, Ahmed Touati, Jesse Farebrother, Mateusz Guzek, Anssi Kanervisto, Yingchen Xu, Alessandro Lazaric, and Matteo Pirotta. Zero-shot whole-body humanoid control via behavioral foundation models.
Ahmed Touati and Yann Ollivier. Learning one representation optimize all rewards. Advances in Neural Information Processing Systems, 34:13-23, 2021.
Ahmed Touati, Jérémy Rapin, and Yann Ollivier: Does zero-shot reinforcement learning exist? arXiv preprint arXiv:2209. 14935, 2022.
Masatoshi Uehara, Issei Sato, Masahiro Suzuki, Kotaro Nakayama, and Yutaka Matsuo. Generative adversarial nets from a density ratio estimation perspective. arXiv preprint arXiv:1610.02920, 2016.
Núria Armengol Urpí, Marin Vlastelica, Georg Martius, and Stelian Coros. Epistemically-gudeed exploration. arXiv preprint arXiv:2507.0544 2025.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL bttps::/proceeet wttp:eeeeeeiinnn...nudeeeee www.ips.cc/pppppppeppp httpauppps.cppppeer_iieeee//e//e www.fflleppppee fttps///ppppe//11/// fale/3f5ee26e 1 04 Papeeaapprreee .upper.pff.. .ddf.
7dee 91
Yifan Wu, George Tucker, and Ofir Nachum. The laplacian in RL: Learning representations with
efficient approximations. In International Conference on Learning Representations, 2019. URL
https: http///pp/eeeiieeee/eee/ee ew net htt.aanrrreeeeeeefooomiieeeeeeee
Jinwei Xing, Takashi Nagata, Kexin Chen, Xinyun Zou, Emre Neftei, and Jeffrey L Krichmar. Domain adaptation in reinforcement learning via latent unified state representation. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 10452-10459, 2021.
13

---

Published as a conference paper at ICLR 2026
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.10742 2020.
Chuning Zhu, Xinqi Wang, Tyler Han, Simon Shaolei Du, and Abhishek Gupta. Distributional successor features enable zero-shot policy optimization. In The Thirty-eighuh Annual Conference on Neural Information Processing Systems, 2024. URL https: https////ppeeeeeeeeeeeeee net forum? id-8 Iyamg2te4.
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta-leamning 2020. URL https: arxiv. org/ abs 1910. 08348.
Ilya Zisman, Vladislav Kurenkov, Alexander Nikulin, Viacheslav Sinii, and Sergey Kolesnikov.
Emergence of in-context reinforcement learning from noise distillation. In Forty-first International Conference on Machine Learning, 2024. URL https: /tppr//e///eeeewewe//e/e net Forum?id=
Y8KsHTIKTV.
Ilya Zisman, Alexander Nikulin, Viacheslav Sinii, Denis Tarasov, Nikita Lyubaykin, Andrei Pol- ubarov, Igor Kiselev, and Vladislav Kurenkov. N-gram induction heads for in-context rl: Improving stability and reducing data needs, 2025. URL https / www.eiiwww....a.eeeeeeeeooeeeeee /411. 01958.
14

---

Published as a conference paper at ICLR 2026
A EXTENDED RELATED WORKS AND BACKGROUND
A.1 BACKGROUND
Contexual Markov Decision Process. Throughout paper we will be dealing with a Contextual Markov Decision Process (CMDP), defined by a tuple MM) where C is a context space and S, A are shared state and action spaces across environments. Function M maps particular context C. to respective MDP, i.e, M(e) A, T Re, ,,,,, ,,,7) with context-dependent transition function T 8 X X C + S. μe being an initial distribution over states and Y E (0,1) discount factor. Intuitively, the context c E C represents a fixed environmental configuration, such as obstacle positions, layout geometry, dynamics vector parameters seed. Throughout this work, the context remains static within each episode, consistent with prior literature (Modi al., 2018; Kirk et al., 2023; Teoh al., 2025). A policy : S + AA is optimal for context e for the reward function R if it maximizes expected discounted future reward, i.e., TT, T2,n( an) =
arg m E E R(s; de)]so, 6o, n, cl.
When the context is fully observable, augmenting the state space with the given context reduces the CMDP to a standard MDP, eliminating the need to model distinct dynamics T rewards It or initial states f However, if the context is partially observable, the learned model must infer and track the uncertainty over true hidden configuration to maintain theoretical optimality guarantees. Such task can be framed as posterior estimation p(c(H) or belief over possible contexts given accumulated history H.
Most successful methods for deriving an optimal policy across arbitrary tasks from a fask-agnostic dataset leverage successor features (Dayan, 1993; Barreto al., 2017; Borsa al., 2019; Park et al., 2024; Zhu et al., 2024) or their continuous counterpart, successor measures (Blier et al., 2021; Touati & Ollivier, 2021; Touati et al., 2022; Agarwal et al., 2025; Jeen et al., 2024). In this work, we focus on the latter framework, specifically its instantiation via backwad representations (Touati & Ollivier, 2021). Below, we briefly outline its key properties.
Zero-Shot RL. Given an offline dataset of transitions D = {(x;, a; 81+1) i= generated by an unknown behavior policies, the agent's objective is to learn compact abstraction of the environment from which it is possible. At test time, this abstraction helps to obtain optimal policy for any reward function Ftest which defines a particular task. Reward function can be specified either dataset of reward-labeled states Drest {(s,, Trest [s:)] or as a direct mapping - (s)s While some prior works assume access to the context labels (Gregor et al., 2019), we focus on the setting where the context is unknown and must be inferred from the data. Alternative formulations of zero-shot RL exist under other formalisms, and we refer to (Kirk et al., 2023) for comprehensive overview.
small
A.2 RELATED LITERATURE
Domain Adaptation and Transfer Learning in RL. While our work will focus on domain adapta- tion applied to estimating successor measure for various dynamies mismatches, we start by briefly reviewing more general ideas in classic domain adaptation and refer to (Kouw & Loog, 2019) for detailed overview. Most methods for domain adaptation can be categorized into importance-weighting (Bickel et al., 2007; Uehara et al., 20 Sønderby et al., 2016) and (Fernando et al., 2013; Eysenbach et al., 2021; Xing et al., 2021; Zhang et al., 2020) approaches. Former methods estimate the likelihood ratio of examples under samples from target domain versus samples from source. which is then used to recalibrate examples from the source domain. The latter approaches learn a unified representation of the environment, targeting to extract only task-relevant abstraction, negating distracting information.
feature learning
The most relevant approach which enables FB representations to generalize across dynamics is Contexual FB (Jeen & Cullen, 2024). This approach uses importance-weighting formalism and introduces two classifiers, which estimate the likelihood of transitions (st, az) and (st, a,, 5t+ being from train or test context and augment the reward function to account for those discrepancies in the dynamics. If augmented reward function lies in the linear span of the Z space during FB training, then the policy can be extracted as described in Equation 3. However, such an approach requires training classifiers from scratch for each novel layout of the environment, limiting its applicability.
15

---

Published as a conference paper at ICLR 2026
Meta-RL. Another major line of related works, Meta-Reinforremen Learning (Meta-RL), focuses on few-shot domain adaptation to unseen tasks or dynamics (Beck et al., 2024). The significant part of research in Meta-RL is dedicated explicitly learning the belief by collecting history of interactions with the environment on inference during test-time (Zintgraf et al., 2020; Dorfman et al., 2021; Rakelly et al., 2019). However, recent works show that it is possible to quantify the belief without learning the posterior implicitly (Laskin et al., 2022; Lee et al., 2023; Zisman et al., 2024; Sinii et al., 2024; Zisman al.. 2025; Tarasov et al., 2025; Polubarov et al., 2025). Leveraging in-context ability of transformers Vaswani et al. (2017), one can learn an end-to-end supervised model, while the transformer's 's context will absorb into robust representation the adaptation- rellvvant information thus enabling fast adaptation. We also leverage this in-context ability to construct the belief representation of the dynamies the agent currently in, but instead operating zero-shot manner.
B PROOFS
Notation recap. Let M MM,s, a,-) be the successor measure of policy π and the reference state- action measure used by FB training. As in the main text, FB seeks low-rank factors F, B such that
M as
F(s, a, 2) B(s a') p(ds'da')
for policies क = Nz. For set of k CMDPs with optimal policies {mt) and successor measures {M*) Mi we define the worst-case class approximation error
o == inf max F.B Isisk
M - F(.,,2)))))) B(.) (P(P)
We write F, B for the trained factors and set the (finite-sample/ /otpsammppppp/tttiioonnn optimization) training discrepancy
BB IEE
Aest.
max ISiSA
where (F*,"") is a minimizer in the definition of et (any minimizer will do). Unless otherwise test distribution Prest that is absolutely continuous w.r.t. p
noted evaluate expectations (Assumption in the main paper), with density ratio bounded by = sups,a dp
dpeset (s, a) < 20.
Lemma (Uniform successor-to-vald stability). Suppose that for some 2 0,
F(so, an, zn) BB.)- M (No, an; -)
sup (0))90)
S.
p(-)
(P(P)
Then for any bounded reward || S R. Q Q 1 I-n Re.
identities, Q. (su. [r(s, a) MM au, dsda). The linear functional M + fr rd has operator norm < || Combining the uniform (p) error on M/ with the contraction of the Bellman resolvent yields the stated (3/(1- R factor (details in the cited stability proofs; constants unchanged).
Proof sketch. By standard successor-occupaagg
Theorem (Regret bound for multiple dynamics with decoupled errors). Under Assumption I and for any bounded reward || Illoc I the policy extracted from the trained factors for CMDP (namely ne with 2; computed from r and B) satisfies
3 E ((sa))ppo (2 a) Q7 s I-7
R R + Aust).
Moreover, ++ & (monotonicity k)).
Proof. Applying Lemma with et + Aest yields || Q, | S 2 R(et + Aaa). Taking
expectation gives the displayed inequality since E [f] < || Monotonicity is immediate because over larger index set cannot decrease.
n
16

---

Published as a conference paper at ICLR 2026
Discussion (Theorem 1). The upper bound separates intrinsic model-class term Ej (harder when more heterogeneous CMDPs are included) from a finite- cate--aapppiiiizzationn term Dest (which can shrink with more data). Thus, adding CMDPs enlarges the worst-case approximation class but may still reduce empirical regret if est decreases.
Assumption 2 (Block-separable parameterization). There exists partition (5₁) of task directions and routing function g 2 + [L] such that the model uses disjoint parameter blocks (F), By): for E Sj the prediction is (s a, 2)) (-(+) and other block parameters are used.
Theorem 4 (Decoupling under block-separable parameters), Assume Assumption 2. Let kmax = max; Sy Then the training objective decouples across blocks j. and the worst-case uniform class error satisfies
max sis,1 S &Ekmux
ISJSL
Consequently, the regret bound in Theorem depends on kmax (not on k).
Proof (sketch). By Assumption 2, losses from tasks = E Sj depend only on (Fj, B2), hence the empirical and population objectives decompose as a sum ΣØ=1 (F By). Minimizers are obtained by solving each block independently. The definition of em the optimal uniform (p) error m tasks then yields et max, fis S E
Discussion (Theorem 2). Partitioning into disjoint cones removes interference: optimization decouples by block, so adding new cones does not inflate the worst-case error beyond the hardest block. Practically, once F, B have enough capacity for the largest block (a > kmax in tabular analogy), the class error can be driven to zero without growing with k.
Let (M,;) } be collection of successor measure of the optimal policies {mi for distinet CMDPs. Given a reference measure on S x A, define the worst-case class approximation error as
= inf max ||Mr, ((, F,B <
(8)
B. FORWARD- BACKWARD (FB) TRAINING
Successor measure and FB factorization. For a policy π and discount y E (0,1), the successor measure M M ap. .) is the (discounted) future occupancy of next states,
M (So, doy X)
Σ 7' Pr((141
80, T)
X XCS,
and, for state-based rewards r S.
²¹ (sp. an) = / ((st) (so, ag ds+).
FB approximates Rd (hence all (T)) with finite-rank factorization conditioned on task vector
E C
(s, ds+) = (F(s, 2), B(s+)) p(ds+),
where F: S Ax Rd is the forward map. B : + Rd the backward map, (..) denotes the Euclidean inner product, and p a reference distribution over next states drawn from the offline dataset. From the factorization it follows that
(F(s, a((s,,,=), 2), B(s+)) p(dst) O7 ~ )
= (F(8, 2x); 2, A Egtro (r(t
Greedy policy family. For each = Sd-1 the greedy policy associated with the representation is
(s) arg aEA
max
(F(s, a, 2),2). 2).
(10)
In discrete action spaces we take the exact maximizer; in continuous control to approximate equation 10 (DDPG-style).
actor network
In some variants B depends (s, a); implementation B(s) in the original formulation.
17

---

Published as a conference paper at ICLR 2026
Bellman identity for the successor measure. Let st+1 T(- st, at) and a++1 Tz(. (: 81+2). For any anchor s* ~ p, the successor measure satisfies
T (St, ag, ds*t) 1(st SL+1) + R M ($141, 444,,
p(ds+)
p(ds*)
(11)
"density" w.r.l.
FB enforces this identity by regressing the scalar score (F()), B(*))) against the right-hand side across random anchors s+.
Training objective (anchor regression). Given dataset D ((st, ot; 8141)], sample 2 2 (e.g.. uniformly on $ or from a mixture that also uses B), compute a++1 R (8(+1) via equation 10, and draw anchors s+ p. Using target networks F,, (Polyak-averaged), the FB loss
FF
E B( BEnz Egtrp
B(s))) 11 =
(F (F(s11, 4411 2), B(5
(12)
On discrete replay buffer (finite p), expanding the square in equation 12 yields the practically convenient equivalent form
Cra Re 222 ((F 01; 2), B(s)) (1))), B(s+)))²
(13)
2 (F ae 2), B(8c+1))
which we use implementation. Gradients update (F, B) while F, B are updated by slow averaging. The actor (continuous actions) is trained to maximize + (F(s, a, 2).
Zero-shot RL procedure (test-time). FB trained without rewards in an unsupervised regime. At test time, for new task specified by reward function r (or small set of labeled states (s:))), we:
1. Infer the task vector. Form
2r Eggt [r(st)r((tt))t))) B(s+)]
2. Act greedily wrr.. Zy. arg max(F (F( 2r),
Use the policy Te in equation 10: (s)
2r lies (approximately) in the linear span of task vectors encountered during training, then O (s, x (F(s, a, 2,),), zn) and uæ, is near-greedy for Qr in the sense of our analysis.
Practical notes. (i) We normalize z, Bss) to the hypersphere for stability; (ii) we mix uniform and Backward-inducee sampling for during training; (iii) target networks and large anchor batches stabilize the regression in equation 12 12-equation 13; (iv) in continuous control we learn an actor (DDPG-style) to approximate the argmax in equation 10. The entire pipeline requires reward labels during training, enabling zero-shot extraction for arbitrary test-time rewards.
C ENVIRONMENT DESCRIPTIONS
C.I RANDOMIZED-DOORE
The Randomized- -ooos MiniGrid environment (Figure 8) a discrete-state, discrete-action finite horizon deterministic environment in which agent has an objective to go to goal location with maximum return 1. Each episode terminates after 100 steps or after reaching goal location. The randomization determines possible open doors locations, fully specifying particular layout. In our experiments, the observation state of an agent consists of (æ, y) coordinates tuple, making it partially observable. Such setting requires properly update beliefs over unobservable layout configuration type. The action space consists of four actions, namely (up, down. right, left), ), while (æ, y) coordinates across both axes are bounded by grid size, which take to be x 9.
18

---

Published as a conference paper at ICLR 2026
(a) First type
(b) Second Type
(c) Third Type
Figure 8: Several possible layouts are visualized, each corresponding to unique possible doors configurations. The agent is denoted as a red triangle. The task specification (goal position) with reward of is denoted by green square and is also randomized. It is custom implementation based on Empty MiniGrid (https:/ //mm/nnn mininiiiii id. far ama org).
C.2 RANDOMIZED FOUR-ROOMS
The Randomized Four-Rooms MiniGrid environment Figure 9 is modification of classic Four- Rooms and is discrete-state, discrete-action, deterministic partially observable environment. For each episode, the maze layout (grid type) is generated randomly, ensuring all the four rooms are connected with exactly single door. Observation state consists of (x, y) coordinates, making this environment hard and checks whether agent could successfully estimate uncertainty over hidden configurations solely based on number of occurrence of each transition, recovering dynamics. In our experiments, we consider X 11 bounds for height and width.
Observation space consists of raw discrete (x, v) coordinates on the grid, while actions correspond to a set possible moves (up, down, left, right). For every layout we record 500 episodes of length 100, yielding dataset that covers almost all possible (s, a) transitions. For testing unseen configurations, we fix agent starting position to coordinates the first empty cell and evaluate performance across 3 static goal positions, farhest away from starting position.
Figure 9: Different layout configurations from randomized Four-Rooms environment. During inference, the goal for the agent (depicted in blue) is to achieve green location. In our experiments we fix starting agent position and fix 3 goals, one for each room.
C.3 ANT-WINS WIND
The AntWind environment modified version of the Ant locomotion task from the MuJoCo simulator, commonly used to test an agent's 's adaptability to changing dynamics. In this environment, an ant-like robot must learn to move forward while being subjected external wind forces varying in magnitude and direction. In experiments consider 17 environments for training, covering three quadrants of possible wind directions on the circle, while leaving others for test, checking extrapolation the fourth quadrant.
For our experiment, we collect dataset by training SAC (Haarnoja et al., 2018) on 3/4 of all possible directions, which results in 16 environments and hold out the other 1/4 for evaluation. Resulting dataset consists of 3400 transition tuples, where each environment configuration is represented as trajectory of length 256.
19

---

Published as a conference paper at ICLR 2026
C.4 RANDOMIZED POINTMASS
Randomized Pointmass a modification of pointmass environment from D4RL Fu et al. (2020). Each episode the environment grid structure is randomized, ensuring all cells are interconnected. The observation space consists of (z, y) transitions. Start position is determined as a first empty cell, while goal location is chosen to be the fartherst away from start (based on Manhattan distance) and ensuring existence of at least one valid trajectory (e.g., through BFS).
Observation space consists of (global 2, global y) position, similar Four-Rooms. We fix dataset size to be le,, only varying number of layouts and episodes per layout, while fixing episode length to 250. We use explore policy, which is random policy with portion of actions repeated (" ("sticky-attions".
Figure 10: Examples of pointmass grid variations.
D EXPERIMENTS DETAILS
Randomized-Doors, For didactic example from Section 3.1 we collect diverse dataset from different layout configurations (open door locations) such that visitation distribution over all states is non-zero. Black color denotes obstacles. The episode length is set to be 100, which is equal to the context length of the transformer encoder for this experiment. Overall, we collect 500 episodes per layout and coverage heatmap is visualized in Figure 11.
Table Comparison proposed approaches against baselines on test (unseen) environments. Results for Fourrooms and Pointmass are averaged across 20 mazes configurations
Method
Environment (Test)
Random
Vanilla-FB
HILP
Lap
Belief-FB
Rotation-FB
Randomized- Randomized-Fourroons
0.05 #0.01
0.15 +0.00
0.2 #0.02
0.1
0.4
0.61 +0.02
Randomized- Pointmass
0.03 10.00
0.1 10011
0.25 20.02
0.1 10.
0.45
20.05
0.55
10.05
Ant-Wind
250 2000000
250
410
1:40.5
290
222.5
550 250.5
640 2:007
Table 3: Comparison of proposed approaches against baselines on train environments. Results for Fourrooms and Pointmass are averaged across 20 mazes configurations.
Method
Environment (Train)
Random
Vanilla-FB
HILP
Lap
Belief-FB
Rotation-FF
Randomizze-FFurroomn
0.18 #0.02
0.25 #0.02
0.4 #0.02
0.2 #0.1
0.7 =0.02
0.85
=0.02
Randomized-Pointnn
0.0 20.05
0.2 #0.2
0.45 #0.1
0.15
20.15
0.76
10.18
0.88 #0.2
Ant-Wind
190
390
410
340
1
680
740 2:70
Note on relience on random exploration during test time. Random exploration relience of BFB and RFB in highly complex environments may fail to discover crucial states needed to disambiguate dynamics identification. However, we emphasize that our work addresses a distinct bottleneck: existing behavioral foundation models (BFMs), particularly FB, tend to collapse when trained on offline data composed of mixed CMDPs. Consequently, training BFMs on large scale mixed multi-modal (in terms of dynamics) data would yield an averaged policy, thus limiting their current
20

---

Published as a conference paper at ICLR 2026
(a) Randomized-
(h) Randomized Four-rooms
Figure 11: Empirical state occupancy measures (p) visualizations of collected datasets for discrete-
based environments.
Vanilla-FB
BFB
RFB
Train
Test
Figure 12: Q-funetion and deterministic policy visualizations (Equation 3) on Randomized Four-Rooms environment. Vanilla-FB ignores environment structure and resulting policy moves through obstacles. BFB and RFB do not such issue.
applicability to unimodal datasets (in terms of dynamics mismatch). Both BFB and RFB overcome this collapse. Developing smarter test-time exploration strategies to streamline dynamics identification remains an important direction for future research.
D. DATASET GENERATION
For Randomized Four-Rooms, produce four training datasets with the following parameters:
21

---

Published as a conference paper at ICLR 2026
episodes per layout
episode length
# Transitions
# layouts
1000000
10
1000
100
1000000
20
500
100
1000000
30
250
100
1000000
50
150
100
Table Details for Randomized Four-Rooms datasets
Randomized Four-Rooms. For experiments on Randomized Four-Rooms during dataset collection we generate randomly grid layout, ensuring that each room is interconnected by exactly one door. For evalution we fix agent start position to (1, 1) with the goal of reaching 3 other goals, specified at other rooms. Each episode terminates after 100 steps. The evaluation protocol is averaged success rate across 3 across 20 environments.
AntWind. For AntWind we first collect trajectories by varying wind direction and training an expert-like SAC agent. After training, collected evaluation trajectories from trained agent. This ensures that all directions are covered and explicitly sets dynamics context, As said in Experiments section, train 16 environments with wind directions corresponding first 3 quadrants of circle, leaving other 4 (last quadrant) for hold out.
E IMPLEMENTATION DETAILS
E. FORWARD-BACKWARD REPRESENTATIONS
E.1.1 GPUs
We run each experiment on Nvidia RTX 4090. The overall training time (for both dynamics encoder and FB training) is approximately hour.
1.2 ARCHITECTURE
The forward-backwartt architecture described below mostly follows the implementation by Touati et al. (2022). All other additional hyperparameters for BFB and RFB are reported in Table 5. Moreover, we should emphasize that our choice of transformer architecture for fayn is mainly based on its abilities to encode large sequences, and other architectural designs (e.g State-Space Models, RNNs) can also be used. This choice does not change observations from Section 3.2, Section 3.3.
Forward Representation F(s, a, 2). The input to the forward representation F is always prepro- cessed. State-action pairs (s,a) a) and state-task pairs (s, 2) have their own preprocessors which are feedforward MLPs that embed their inputs into a 5 512-dimensional space. These embeddings are con- catenated and passed through third feedforward MLP F which outputs d-dimensional embedding vector. Note: the forward representation F is identical to & used by USF so their implementations identical (see Table 5). Also, for stability reasons of TD learning, we make ensemble of F and take their mean as aggregation function.
Backward Representation B(s). The backward representation a feedforward MLP that takes state input and outputs d-dimensional embedding vector.
Actor (ss 2). Like the forward representation, the inputs the policy network are similarly preprocessed. State-action pairs (s, a) and state-task pairs (s, have their own preprocessors which feedforward MI.Ps that embed their inputs into 512-dimensional space. These embeddings are concatenated and passed through third feedforward MLP which outputs a-dimensional vector, where a is the action-space dimensionality. A Tanh activation used on the last layer to normalise their scale. Note the actors used by FB and USFs are identical (see Table 5). For discrete environments, optimal policy is greedy, while for continuous DDPG-style is used for approximating a argmax.
Misc. Layer normalisation and Tanh activations used in the first layer of all MLPs to standardise the inputs as recommended in original paper for both discrete and continuous becnhmarks. Baseline taken from official repository contrall lable agent.
22

---

Published as a conference paper at ICLR 2026
Table Hyperparameters for FB. Hyperparameters for Belicf-FB and Rotation-FB are highlighted
in
Hyperparameter
Value
Latent dimension d
150 (100 for discrete)
/ v dimensions
(1024, 1024) (256, 256, 256) (1024, 1024)
B I yp dimensions
Preprocessor dimensions
Std. deviation for policy smoothing
0.2
Truncation level for policy smoothing
0.3
Learning steps
1,000,000
Batch size
1024
Optimiser
Adam
Learning rate
0.0001
Learning rate fayn
0.0001
Discount "
0.99, 0.98 (Maze)
Activations (unless otherwise stated)
GeLU
Target network Polyak smoothing coefficient
0.05
z-inference labels
10,000
mixing ratio
0.5
K
50. 100 for Pointmass
Contexual representation h dimension
150 (100 for discrete)
Next state predictor Spred
(256, 256, 256)
E.2 HILP
We take official implementation in JAX from Park al. (2024), All of the hyperparameters are the same as in original paper.
E.3 TASK SAMPLING DISTRIBUTION Z
Vanilla-FR. FB representations require a method for sampling the task vector z at each learning step.
Touati al. (2022) employ mix of two methods, which replicate:
1. Uniform sampling of 2 on the hypersphere surface of radius vd around the origin of Rd
2. Biased sampling of by passing states & ~ through the backward representation = Bss). This also yields vectors on the hypersphere surface due the L2 normalization described above, but the distribution is non-uniform.
We sample ~ 50 50 (either randomly or from B3) from these methods at each learning step in original work by Touati & Ollivier (2021).
Rotation-FB. After transformer Jayn pretraining stage, RFB at each gradient step chooses task- conditioning vector 2F8 based on i) context representation acting as axes coming from faya and ii) drawing task encoding vectors 2p around this axes. We also perform normalization as in Vanilla-FB by projecting resulting vector surface of hypersphere of radius Vd.
Stage ii) is implemented as drawing samples as Рнв ~ vMF(μ = h, k).. In order to remove high computational costs, implement this sampling procedure through Householder reflection around context axes, by first drawing z from one of the basis vectors (e.g., north pole) and then performing rotation.
23

---

Published as a conference paper at ICLR 2026
E.4 PSEUDOCODE
Algorithm 1 Belief-FB Training
1: Input: offline diverse dataset D consisting of transitions based on hidden configuration variable
2: Initialize transformer encoder fayng Fg. Bu, number of gradient steps for transformer pre-training K,
context length 7, Polyak cocfficient, 3, batch size B learning rates A Ar, AB
do 4: sample batch of B trajectories of length {(8,,t, Gi (1,4+1)]i T. ~ D actt, $1,141 =1, I,..,, B,
3: while update steps
(yui log o)), ; fayan(1 (1 (s,,t,
6: 2; p. E; exp(poog oi:).
Z. Zdyn,,
I.....T &eepresentation zaya is shared across each sequence
8 = (prea(S1.e as,e, Zr.e) == 1,...,77 T, 1 1,...,
9: Lionee:: BT
Ei |Lee11 81,411 Vo Conteu
10:
0sayn faayn
end while
12: while not converged do
Af V (TF 14: f- wB Aa (F
13:
# #BB training, Equation ??
15:
end while
Algorithm 2 Sampling 2rB for RFB
Input: B (batch size), (latent dimension), anchor matrix H R Bxd, (concentration)
Output: ZERBXX
|l2 E) 2: S VMF SAMPLE NORTHPOLE(B,B.... k)
1: Normalize anchors: u, 4 H.
for i = I,... B
> draw B VMF samples
3: for i + to B do
4: R. HOUSHHOLDER HOUSSHOOLLRR_LOTTTTOOii)
5: 2; <- R
6: end for
7: Z t
((z;)/
8: return
24


Reviews:
[{"review_id": "1haPxZpMc8", "reviewer": "Reviewer_EvRU", "summary": "This paper proposes Belief-Forward-Backward (BFB) and Rotational Forward-Backward (RFB), two methods that extend Behavioral Foundation Models (BFMs) for zero-shot adaptation to unseen dynamics.  \n1. Belief-Forward-Backward (BFB): introduces a transformer-based belief encoder that infers environment dynamics and conditions the Forward-Backward (FB) representation on this latent belief, improving adaptability to new transitions.  \n2. Rotational Forward-Backward (RFB): further refines BFB by aligning latent policy directions with context-specific clusters using a von Mises-Fisher prior, allowing better separation between dynamics in the latent space.  \n3. The paper provides theoretical explanations for why standard FB representations fail under varying dynamics and offers formal regret bounds for BFB and RFB.  \n4. Experiments across discrete (Four-Rooms) and continuous (PointMass, AntWind) domains demonstrate improved zero-shot generalization of BFB and RFB compared to baselines such as FB, HILP, and LAP.", "questions": "Empirically, RFB consistently outperforms BFB. Are there situations where the authors believe BFB would be preferable to RFB—for instance, when dynamics clusters are fuzzy, data is limited, or tuning κ is undesirable?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "1. The paper presents a clear and well-motivated idea that is straightforward to implement.\n2. The visualization and qualitative analyses are helpful for understanding the model behavior.\n3. Experiments convincingly demonstrate the effectiveness of the proposed methods across multiple settings.", "weaknesses": "1. **Minor issues:**  \n   - Equation (1) uses `(s_{t+1}, a_{t+1}) ∈ X` although `X` was defined as a subset of states. It should be `s_{t+1} ∈ X`.  \n   - Line 144 appears incomplete.  \n   - Theorem 1 defines `‖r‖_∞ ≤ R` but the symbol `R` is not used in the subsequent expressions.\n2. The paper lacks a formally defined problem statement. The current description is mostly natural language; a mathematical setup of the CMDP and belief inference process would improve clarity.\n3. Figure 2 is conceptually confusing without such a formal section. It appears that observations are included in the state, and layouts (contexts) are not distinguished in state space, meaning the agent cannot infer context. Since `z_FB` does not encode context either, the result—an averaged policy—becomes unsurprising.\n4. The connection between Theorem 1 and the statement in line 215 is logically weak. Theorem 1 implies that more CMDPs increase the *upper bound* of the worst-case error, not necessarily the regret itself. More CMDPs may also mean more samples, potentially reducing `ε_k`. Section 3.1 lacks discussion of this, reducing its persuasiveness.\n5. The value of Theorem 2 depends on the validity of Theorem 1, and it has the same limitations. In addition, the variable `L` is undefined in Theorem 2.\n6. The experimental comparisons may be unfair: BFB and RFB explicitly encode dynamic information into `z_FB`, while baselines like FB or HILP do not. This makes it expected that those baselines perform close to random.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2QtQs8ISyM", "reviewer": "Reviewer_jHoj", "summary": "Behavior foundation models (BFMs), such as Forward-Backward representation, enable zero-shot policy optimization under test-time rewards. An example is Forward-Backward representation, which represents the success measure as a product of forward features $F(s, a, z)$ and backward features $B(s^+)$, and extracts policies as $\\pi_z := \\arg\\max_a F(s, a, z)^\\top z$. However, FB inherently encodes information about the dynamics and doesn't transfer to unseen dynamics. This paper first diagnoses this issue by visualizing the optimal actions corresponding to latent codes and find that training on data from different dynamics results in interference. To address this, they propose Belief-FB (BFB), which introduces a context encoder -- a permutation-invariant transformer that encodes a set of transitions into a context vector about the dynamics. The context $h$ is concatenated to the latent code $z$ and used to condition the forward representation $F(s, a, z+h)$. However, because FB training draws latent $z$ uniformly, this still entangles all the policies from different dynamics. To further disentangle the polices, the paper proposes Rotation-FB (RFB), which samples $z$ within a cone around a projection of the context vector $h$. The paper evaluates BFB and RFB in discrete and continuous state-based environments and finds BFB to outperform baseline BFMs without contextual awareness, and RFB to further improve performance. Overall, the paper takes an important step towards removing the stationary-dynamics assumption of BFMs.", "questions": "**Major**\n1. Why can't we simply add context to the observation? For example, one can encode the maze layout and concatenate it with the rest of the state vector (or even better, use image-based policies). This wouldn't work when the context is partially observable, but for fully observable environments, can you evaluate this simple oracle baseline?\n2. Can you add comparison to immediately relevant baselines that are dynamics aware, such as [1]?\n3. Can you demonstrate the effectiveness of your method on more challenging domains, e.g., dexterous manipulation, where the friction parameter is varied?\n\n**Minor**\n\n4. Typo line 342: \"platoes\" -> \"pleateaus\"\n5. Typo line 732: missing right bracket\n6. Lines 119 - 132 are not very clear for readers unfamiliar with the FB literature. For example, the policy is defined as argmax (saying it is extracted as argmax might confuse readers). I suggest adding a section on the arithmetic of FB in the appendix. \n\nReferences:\n\n[1] Scott Jeen and Jonathan Cullen. Dynamics generalisation with behaviour foundation models. Workshop on Training Agents with Foundation Models at RLC 2024.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The proposed methods are intuitive. Introducing a context vector to address ambiguity during FB training across diverse dynamics is a logical and effective approach, while sampling $z$ around a cone-shaped prior makes sense as an inductive bias.\n2. The visualizations (Figures 3 and 4) are highly informative, providing compelling evidence for the policy entanglement issue in standard FB and demonstrating the improvements achieved by BFB and RFB.\n3. The paper provides a rigorous formal analysis, showing that (1) the worst-case regret for vanilla FB increases with the number of environments and (2) the FB approximation error is reduced by introducing the cone sampling approach.\n4. Both BFB and RFB outperform the baseline methods in the empirical evaluations.", "weaknesses": "1. The evaluation environments are excessively simple.\n2. The paper does not compare the proposed methods to immediately relevant baseline with context awareness (e.g. contextual FB [1]).\n3. The vanilla context vector doesn't fully disentangle policies, which is why the paper introduces additional rotation variant.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uI0gM2y8Se", "reviewer": "Reviewer_mdha", "summary": "This paper identifies a critical failure in Forward-Backward (FB) based Behavioral Foundation Models: training on mixed-dynamics data causes interference, as the model averages the successor measures across all environments. This averaging entangles the latent policy representations, yielding sub-optimal policies. To solve this, the authors propose Belief-FB (BFB), which introduces a transformer-based belief estimator to infer a dynamics-specific context vector from reward-free transitions, which then conditions the FB model. They further introduce Rotation-FB (RFB), which structures the latent space by sampling from a von Mises-Fisher (vMF) distribution centered at the inferred context. This alignment partitions the policy space into non-overlapping, dynamics-specific clusters, preventing interference and enabling zero-shot adaptation to unseen dynamics. The methods were tested on simulated discrete and continuous tasks with changing dynamics, including Randomized Four-Rooms, PointMass, and AntWind. The results show that both BFB and RFB successfully adapt to seen dynamics and generalize to unseen ones, achieving up to 2x higher zero-shot returns compared to baselines like vanilla FB.", "questions": "Lines 143-144. Something appears to be missing?\n\n### References\n\nPeng, Xue Bin, et al. \"Learning agile robotic locomotion skills by imitating animals.\" arXiv preprint arXiv:2004.00784 (2020).", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The authors identify a fundamental failure mode in Forward-Backward (FB) models, which is termed interference.\nThe proposed method demonstrates promising zero-shot generalization improvement to unseen dynamics.", "weaknesses": "The work addresses a problem specific to the FB paradigm and builds upon this methodology only. It may not be general to other behavioral foundation models.\n\nThe introduction of a context into a dynamics model to disambiguate between distinct training environment dynamics is a well-established approach in adaptive control and reinforcement learning (Peng et al.). However, I do not see a discussion of these previous works (or really any related works discussion).\n\nEvaluations are in simulated, ideal environments under narrow dynamics variations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JG1u0efL5W", "reviewer": "Reviewer_dycu", "summary": "The paper introduces an algorithm for unsupervised zero-shot reinforcement learning under partial observability. It considers environments where the dynamics are governed by a latent parameter that is not directly observable but can be inferred from transitions. The authors propose using a permutation-invariant, transformer-based architecture to build an estimate (or belief state) of the dynamics from a set of transitions. This belief vector is then used to condition the Forward-Backward model, i.e., both the actor and critic are trained conditioned on the belief state. Furthermore, they propose modifying the structure of the task/policy space to be belief-state dependent, resulting in clusters of dynamics. The effectiveness of the method is evaluated on three toy domains.", "questions": "- I found Section 3.1 a bit trivial. It is expected that a method designed for a single MDP does not perform well in the considered contextual MDP. What am I missing? Why spend so much time justifying this aspect?\n- I like the approach of building the latent belief by predicting the next state. Could you discuss further the choice of using a non-causal transformer? Why did you use Gaussian regularization?\n- Do you use samples from the training dataset to perform reward inference? Does the performance change if you use transitions from the test environment for inference?\n- You assume access to a set of transitions from the test environment to infer the context. However, you could also compute the context directly while interacting with the environment. Could you evaluate the test performance with respect to the number of samples used for belief inference?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- As far as I know, this is the first paper to tackle the problem of unsupervised reinforcement learning in partially observable contextual Markov Decision Processes (MDPs).\n- The approach is simple and sound, and the authors provide a clear, step-by-step construction of the method, explaining why each component is necessary.\n- Although the core idea is not very novel, I believe its application to reinforcement learning is new and interesting for the community.", "weaknesses": "I don't see any major weaknesses to mention in this section. I will report a few questions in the next section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)