Zero-Shot Policy Transfer with Disentangled Task Representation of Meta-Reinforcement Learning — part 6/6

Cleaned-text code points [72074, 79832). Bibliography omitted here; source.md includes it.

Under review for the Reinforcement Learning Conference (RLC)
Table 4: Online task inference (Thompson Sampling) hyperparameters.
Hyperparameter ColorGrid Wind MetaWorld
σ10 15 † 1
† Shared across domains: Bayesian linear regression with prior precision λ=1, posterior decay 1, exploration bias
κ=0, posterior resampled every environment step.
Table 5: Intrinsic reward hyperparameters.
Hyperparameter ColorGrid Wind MetaWorld
Latent-dynamics-error bonus weight 0.0 0.0 0.0
Episodic novelty bonus weight 1.0 1.0 1.0
Intrinsic critic reward scale 1.0 1.0 1.0
Episodic novelty: nearest neighborsk10 10 10
Episodic novelty: cluster distance threshold8×10 −3 1×10−3 8×10−3
Episodic novelty: maximum kernel similarity 8.0 8.0 8.0
Episodic novelty: kernelϵ1×10 −4 1×10−5 1×10−4
C.4 Intrinsic Rewards
The intrinsic reward at each timestep is a sum of two terms: an episodic novelty bonus computed in
the style of NGU (Badia et al., 2020) on the latent context features, and a term proportional to the
previous-timestep latent-dynamics prediction error (returned by Alg. 1). The combined scalar is
normalized and exposed to the policy as an additional observation, so the actor can directly condition
on the current intrinsic signal in addition to driving theβ-blended critic.
C.5 Actor–Critic Update
Sequence backbone.The actor and the intrinsic–extrinsic critic share a transformer trajectory
encoder with rotary position embeddings (Su et al., 2024) and sliding-window FlashAttention (Dao
et al., 2022).
Actor and Critic.The critic, QR, is a small ensemble of two-hot regression heads over a fixed
return range (Hafner et al., 2023), which helps absorb the large variation in TD-target scale induced
by randomizing α and β. Both the actor and critic heads consume a Fourier encoding of the
(α,β )scalars in Equation 4 as additional input features. The actor is an MLP with a standard
Tanh-Gaussian output.
C.6 Computational Resources and Statistical Significance
All models are trained on an NVIDIA A500 GPU, with training runs for theMetaBFM (Section 4)
experiments lasting 12–36 hours depending on the domain. Shaded error bars in results figures
indicate the mean and standard deviation of at least three trials. Results indicating test-time
adaptation curves (performance vs.k) average these results over the last 3 evaluation cycles during
training, then over seeds.
22

Under review for the Reinforcement Learning Conference (RLC)
Table 6: Actor–critic hyperparameters.
Hyperparameter ColorGrid Wind MetaWorld
Training sequence length 360 226 501
Actor / critic context length 96 128 128
Policy attention window 64 timesteps full attention 64 timesteps
Actor MLP hidden width 512 300 400
Critic MLP hidden width 300 300 400
Critic output bins (two-hot) 64 64 96
Critic return range[rmin,r max] [−400,600] [−50,400] [−300,3000]
Target network rate (τ)5×10 −3 3×10−3 5×10−3
Number of Fourier frequencies for(α,β)encoding 16 8 8
D Limitations
BFM and ZSRL algorithms are an active area of research, and reported gains often depend on
hyperparameters, implementation details, and benchmark choice. Our experiments use TD-JEPA
as a representative state-of-the-art BFM update, but broader comparisons across BFM objectives
and online task inference methods remain important future work.MetaBFM also introduces test-time
choices throughα and β. These parameters are useful because they let a single policy interpolate
between exploration, exploitation, and ZSRL-style reward inference, but they also create a model-
selection problem. In our experiments, we evaluate a few interpretable settings; in deployment, one
may need validation tasks or an online rule for choosingαandβ.
E Additional Results
E.1 Online ZSRL
BFMs are inherently proposed as offline algorithms (Touati et al., 2022; Agarwal et al., 2025b),
trained on offline exploratory datasets, which can then be used to directly infer policies for diverse
reward functions. While the inference involves solving a linear regression problem, as discussed in
A.2, using samples drawn from the same offline datasets. Prior work (Sun et al., 2025) has shown
that these BFMs are not effective for online training. Methods such as (Li et al., 2025; Tirinzoni
et al.) have used expert datasets to guide the online training of these models.
We compare the performance of BFMs when trained offline and online on standard DM Control
tasks (Tassa et al., 2018). We compare (Figure 14) the offline-trained BFM (on ExoRL datasets
(Denis Yarats, 2022)) using offline inference (labeledoffline) with the same model using online
inference (labeledoffline-online) as proposed by (Rupf et al., 2025) and with a fully online BFM
that does not assume access to any external dataset during training or inference.
It is clear that BFMs are optimally suited to fully offline settings. While using online inference
degrades the performance slightly, the performance of online-trained BFM is significantly worse.
BFMs require high coverage and well-explained data, which is hard to obtain during online training,
leading to poor performance. On the other hand, Meta RL often improves when trained online, as
observed from Figure 15.
E.2 Utilizing Available Rewards
BFMs are designed to be trained using reward-free interactions. We design a training procedure
by providing occasional reward information to the BFMs to see whether they can make use of the
23

Under review for the Reinforcement Learning Conference (RLC)
Figure 14:Offline vs Online BFMs:Comparison of fully offline, offline training and online
inference and fully online BFMs on DM Control tasks.
1 2 3 4 5 6
Adaptation Episode
0.0
0.2
0.4
0.6
0.8
1.0Oracle Norm. Return
1 Train Task
1 2 3 4 5 6
Adaptation Episode
10 Train Tasks
1 2 3 4 5 6
Adaptation Episode
100 Train Tasks
1 2 3 4 5 6
Adaptation Episode
10,000 Train Tasks
Train (online)
T est (online)
Train (offline)
T est (offline)
Figure 15:Online vs. Offline Meta-RL.In ColorGrid (n = 5), offline meta-RL from a static
dataset generated by a uniform behavior policy leads to significant distribution shift in long-context
representations; online learning substantially improves performance.
available rewards. On the trajectories where rewards are available, the BFMs exploit the reward
function to interact with the environment using an optimal policy (optimal for the partially trained
BFM). It adds regularization to the state features to ensure they remain sufficiently predictive of the
available rewards, as discussed in Section 3.3.
We compare the performance of online-trained BFMs with and without the use of reward functions
in Figure 16. The improvement in performance when using available rewards is not consistent. While
an improvement is observed in the walker, similar improvements are not observed in the other two
environments. On ColorGrid, the BFMs show improvement in performance when leveraging the
available rewards, as shown in Figure 17.
24

Under review for the Reinforcement Learning Conference (RLC)
Figure 16:Effect of using rewards during training of BFMs:Comparing online BFMs with
and without using available rewards during training.
1 2 3 4 5 6
Adaptation Episode
0.2
0.4
0.6
Oracle Normalized Return
0 50k 100k 150k 200k
Gradient Steps
0.3
0.4
0.5
0.6
Reward R²
0 50k 100k 150k 200k
Gradient Steps
0.00
0.05
0.10
0.15
0.20
zr Pairwise Similarity
Baseline 10% zr 60% zr 60% zr + Reg =1
 60% zr + Reg =10
 60% zr + Reg =100
 60% zr + Reg =1k
Figure 17:Reward Regression andz-sampling.Sampling reward-estimated z values and fitting
known reward labels (using an auxiliary loss with weightλ) fits better linear reward estimates
(center), focuses learning on more similar latent tasks (right), and marginally improves early-episode
adaptation in the ColorGrid domain (left).
25