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

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

Under review for the Reinforcement Learning Conference (RLC)
labeled tasks, even if we cannot cover every behavior needed at deployment. We investigate these
gaps before proposing a solution that enables a competitive hybrid between meta-RL and BFMs.
1 2 3 4 5 6
Adaptation Episode
0.0
0.2
0.4
0.6
Oracle Norm. Return
Seq. RL Baseline
FF LCL  FF RL
Seq. LCL  FF RL
Seq. LCL  Seq. RL
Figure 2:Latent Context Learning. (Left)We learn a timestep embeddinght and sequence
modelSθ(τ:t)that generates an online estimate of the environment parameters,˜ct. The representation
is trained with self-supervised inverse and latent dynamics objectives; Algorithm 1 gives the full
procedure. Action inputsat are necessary context to infer dynamics, but are time-shifted to avoid
leaking the IDM label.(Right)Context features learned by a feedforward (FF) or Transformer (Seq.)
Sθ are presented to either a feedforward or sequence-based supervised RL policyπin ColorGrid.
3.1 Adapting to Novel Dynamics
An RL policy that cannot conditionπ on estimates of the current dynamics,cd, and reward function,
cr, learns an average policy underp(cd,cr)that is generally not optimal for any specific task (Ghosh
et al., 2019). BFMs share this problem, butz∗ already plays the role ofcr, leaving us to infercd.
Following context-based meta-RL (Beck et al., 2023), we augment each trajectoryτ with a learned
estimate ˜cd and trainπ(a|s,˜cd). Because context labels are assumed to be unavailable, we use a
self-supervised dynamics objective. Unlike BeliefFB (Bobrin et al., 2025), we target timestep-level
online task inference with latent dynamics representations that avoid method-specific changes to the
BFM update.
1 2 3 4 5 6
Adaptation Episode
0.0
0.1
0.2
0.3
0.4
0.5Oracle Normalized Return ( )
0.5
0.6
1 2 3 4 5 6
Adaptation Episode
0.00
0.04
0.08
Forward Dynamics MSE ( )
0.00
0.05
0.10
0.15Inverse Dynamics MSE ( )
Baseline LDM + Reg (d=32) LDM (d=32) FDM (d=32) LDM + Reg (d=2) LDM (d=2) FDM (d=2)
Figure 3:LCL and Generalization.TD-JEPA in ColorGridn = 5with various context embedding
sizes and training objectives. Performance correlates with lower FDM and IDM error. LDM
representations recover similar dynamics prediction as FDM without training on this objective.
Figure 2 shows our Latent Context Learning (LCL) architecture, which produces˜cd at every timestep
during rollouts. We embed each timestep of the trajectoryτ:t to create the input sequence{h0,...,h t},
then pass it through a causal sequence model. LetSθ(τ:t)∈R d be the output representation at
timestep t. This output is trained to identify the active dynamics through self-supervised inverse
dynamics modeling (IDM) and latent dynamics modeling (LDM). Because accurate estimates of
cd improve these objectives, and becauseSθ is the only representation with access to the necessary
history, we treat˜ct =Sθ(τ:t)and can train a standard BFM on augmented states¯st = (st,˜ct). Sθ
observes action choices to infer dynamics, but we avoid leaking the inverse dynamics model label by
5

Under review for the Reinforcement Learning Conference (RLC)
staggering the action input by two timesteps. We prefer latent dynamics modeling because forward
dynamics modeling (FDM) scales with the size of the state space. Following latent world modeling
and self-supervised sequence models (Oord et al., 2018; Schwarzer et al., 2020; 2023), we use an EMA
target network forht+1 labels (Grill et al., 2020), and can optionally encourage˜ct to be dissimilar
across sequences in each training batch (Jajoo et al., 2026). Appendix C.1 and Algorithm 1 provide
additional details. By default, we train IDM and LDM jointly; detached FDM prediction error is
used only as a diagnostic metric.
1 2 3 4 5 6
Adaptation Episode
0.0
0.2
0.4
0.6
0.8
1.0Oracle Normalized 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 T est Meta-RL ZSRL (Baseline)
ZSRL w/ LCL Pretrained on 10k T ask Dataset ZSRL w/ LCL Pretrained on Same Dataset ZSRL w/ LCL Trained Jointly
Figure 4:Latent context learning improves BFM generalization.We increase the number
of training(cd,cr)task pairs with separate offline datasets. Standard TD-JEPA collapses across
environments, while LCL-augmented TD-JEPA recovers adaptation and improves test performance
as the training set grows. Pretraining LCL is most helpful at small training set sizes, but joint
training approaches the same performance at scale.
Context estimates should improve supervised RL by reducing or eliminating the need for long-term
memory policies, and Figure 2 (Right) verifies this effect in ColorGrid. By augmenting states with
context estimates, we can also make BFM networks context-aware without the instability of sequence
models. Figure 3 shows that, when the context representationSθ(τ)∈R T×d has enough capacity
to representcd, joint IDM and FDM learns accurate dynamics estimates that improve TD-JEPA
performance. In Figure 4, we define train/test splits of(cd,cr)task pairs with separate offline
datasets, and evaluate generalization as the number of training tasks grows. Since LCL and the
downstream BFM agent are independent, they can be trained separately. Pretraining and freezing
LCL on the largest task set improves BFM generalization at small training set sizes, but gives no clear
advantage over training both systems from scratch on the same offline dataset. This is encouraging
because joint training is the only realistic option when the dataset is collected online.
To further illustrate this BFM limitation and demonstrate the effectiveness of context learning, we
conduct experiments on motion-tracking tasks for G1 humanoid locomotion based on BFM-Zero (Li
et al., 2025). We introduce random dynamics variations through wind and gravity perturbations of
varying strength. Consistent with our analysis above, BFMs take average actions across multiple
dynamics settings, which can lead to performance degradation. In the in-distribution setting, we find
that BFM-Multi performs even worse than BFM-Single under multi-dynamics evaluation, suggesting
that BFMs suffer substantially when they cannot infer the environment dynamics. We evaluate
tracking performance on the LAFAN (Harvey et al., 2020) dataset, which contains4.6hours of motion
sequences. As shown in Table 1, context learning greatly improves performance in the in-distribution
setting, suggesting its promise for cross-domain training. Surprisingly, context learning also helps
to some extent in the out-of-distribution setting. We provide an illustration of BFM failures under
multiple dynamics, along with additional details about the task and environment, in Appendix B.3.
3.2 Exploring, Exploiting, and Learning Online
BFM algorithms are primarily developed on offline RL benchmarks (Laskin et al., 2021) with behavior
datasets generated by an online policy optimizing an intrinsic reward (Burda et al., 2018; Seo et al.,
2021; Raileanu & Rocktäschel, 2020). BFMs rely on high state-action coverage and struggle to learn
6

Under review for the Reinforcement Learning Conference (RLC)
Table 1:BFMs with latent context learning improve humanoid motion tracking.We com-
pare motion-tracking metrics for humanoid locomotion under in-distribution and out-of-distribution
dynamics. BFM-Single trains without wind/gravity randomization and can be viewed as an in-
distribution oracle for BFM-Multi, or as a worst-case baseline under OOD dynamics.
Method In-Distribution Out-of-Distribution
mpjpe-l(↓) vel. dist.(↓) accel. dist.(↓) emd(↓) mpjpe-l(↓) vel. dist.(↓) accel. dist.(↓) emd(↓)
BFM-Single 117.2 1.64 2.75 1.01 194.7 2.81 4.83 1.80
BFM-Multi 132.7 1.77 2.96 1.13 151.6 2.13 3.63 1.39
FDM+BFM-Multi 125.3 1.73 2.90 1.09 149.3 2.09 3.54 1.37
LDM+BFM-Multi121.5 1.68 2.85 1.05 145.9 2.06 3.52 1.36
online because the actor objectiveLactor collapses onto narrow behavior. Appendix E.1 explores
the difficulty of online learning and task inference in a standard locomotion benchmark. DVFB
(Sun et al., 2025) addresses dataset diversity by maximizing a combination ofQSF and a supervised
critic trained on intrinsic curiosity rewards, which we will callQR (Equation 1). Recent applications
of BFMs to humanoid robotics (Tirinzoni et al.; Li et al., 2025) use a similar idea, but deriveQR
rewards from behavior in an expert motion-capture dataset.
LDVFB =−E s∼D,z∼p(z),a∼π(·|s,z)
[
QSF(s,a,z) +αQ R(s,a,z)
]
.(1)
In meta-RL problems, we can take this idea further. First, we replace theglobalintrinsic reward—one
that decreases as a function of total training steps, such as RND (Burda et al., 2018)—with an
episodicreward that decays within ak-episode meta-trial (Henaff et al., 2023). An agent maximizing
this reward is incentivized to explore new states across each episode at test time. We use an
intrinsic reward that combines an off-the-shelf episodic bonus (NGU (Badia et al., 2020) or E3B
(Henaff et al., 2022)) on LCL timestep featuresht (Fig. 2) with the LCL prediction error from the
previous timestep. The latter is a core meta-RL idea (Zintgraf et al., 2021b; Zhang et al., 2021): it
incentivizes exploration that improves the˜cd estimate, which can also improve in-context. Second,
because episodic rewards are non-Markov,QR and π should have memory and take LCL-augmented
trajectories¯τas input. Finally, we sampleαduring training and conditionQR on the active value:
Lα-BFM =−E τ∼D,z∼p(z),
α∼p(α),a∼π(·|¯τ,z,α)
[
(1−α)Q SF(¯s,a,z) +αQ R(¯τ,a,z,α)
]
.(2)
The result is a policy that generates more diverse data during training while allowingα to vary at
test time. Due to our choice of intrinsic reward,α = 1drives the policy to explore new behavior each
episode in hopes of (1) finding rewards that improve OptiBFM’s online task inference ofz∗ and (2)
improving the context estimate˜cd. In contrast,α= 0exploits that knowledge to maximize returns
under z∗. We can scheduleα to recover meta-RL’sk-shot exploration-exploitation tradeoff at test
time, as demonstrated in Figure 5.
3.3 Learning from Available Rewards
BFMs assume access to zero reward labels during training. While this may be necessary when learning
offline, such as from a web dataset that was not intended for RL, the constraint is unnecessary when
learning online in simulated environments. When we create meta-RL or multi-task RL training
domains, we can usually create some reward-labeled tasks; the challenge is creating enough to cover
all behaviors we may want during deployment. Adding training tasks should improve generalization
with diminishing returns, creating a clear path to improving the agent by improving its dataset.
Ideally, we would bring this relationship to BFMs. Each new train-time objective can help span the
space of possible behaviors and supervise learning progress, while unsupervised learning fills in gaps
and finds solutions to objectives we had not considered.
7

Under review for the Reinforcement Learning Conference (RLC)
Adaptation Episode 
1 2 3 4 5 6
Exploit 0 0 0 0 0 0
Explore 1 1 1 1 1 1
Decay 1 .5 0 0 0 0
1-Shot 1 0 0 0 0 0
3-Shot 1 1 1 0 0 0
5-Shot 1 1 1 1 1 0
Random 
1 2 3 4 5 6
Adaptation Episode
0.0
0.2
0.4
0.6
0.8
Oracle Normalized Return
1 2 3 4 5 6
Adaptation Episode
0
5
10
15
20
Unique Cells Visited
1 2 3 4 5 6
Adaptation Episode
0.0
0.2
0.4
0.6
0.8
Oracle Normalized Return
1 2 3 4 5 6
Adaptation Episode
0
5
10
15
20
Unique Cells Visited
 Schedule
 Dense Rewards Sparse Rewards
Performance Exploration Performance Exploration
Figure 5:Exploring at Test Time.ColorGrid (n = 5) policies trained by Eq. 2 can varyα over
k = 6test episodes. High α explores more of the state space, improving online task estimates, which
can then be exploited to outperform the standard BFM policy (α = 0, blue) when rewards are sparse.
We can approach this problem from two directions. The first is to let reward labels improve the BFM
update itself. For example, BFMs learn featuresϕ(s)and assume ϕ(s)⊤z =r; when reward labels
are available during training, we can enforce this relationship directly with an auxiliary loss. It has
also become standard to sample trainingz∼p (z)from recent values ofϕ(s), implicitly assuming
evaluation on sparse state-based goal-reaching (Touati & Ollivier, 2021; Touati et al., 2022). With
reward labels, we can extend this idea to arbitrary reward functions by using online task inference
along the labeled trajectory to estimatez∗, as if the same states and rewards appeared at test time.
Appendix E.2 evaluates both ideas in ColorGrid and DMC ExoRL, finding that they makeϕ more
predictive of rewards but have mixed effects on policy performance. A second approach is to restore
supervised meta-RL where rewards are available and create a true hybrid method. We modify the
“intrinsic” criticQR from Section 3.2 to learn from a weighted combination of intrinsic rewards,ri,
and extrinsic (supervised) rewardsre:
QR(¯τ,a,z,α,β)≈(1−β)r i +βr e/bracehtipupleft/bracehtipdownright/bracehtipdownleft/bracehtipupright
Intrinsic and Extrinsic Reward
+γQR
(
¯τ′,a′,z,α,β
)
(3)
Where ¯τ′ and a′ are the next timestep in the standard TD backup. We are then free to repeat the
idea from Equation 2 of sampling the weightβ during training and conditioning the actor and critic
on the current value so that it may also be adjusted at deployment:
Lαβ-BFM =−E τ∼D,z∼p(z),
α∼p(α),β∼p(β),
a∼π(·|¯τ,z,α,β)
[
(1−α)Q SF(¯s,a,z) +αQ R(¯τ,a,z,α,β)
]
(4)
When β = 0, we recover the exploratory ZSRL objective of Eq. 2, and whenα =β = 0, we recover
standard ZSRL. Most importantly, whenα =β = 1, we recover supervised meta-RL. In practice,
this makes it possible to compete with a standard supervised meta-RL baseline on in-distribution
tasks, while opening the possibility of generalizing viaz to novel tasks when meta-RL fails. Figure 6
visualizes this space of possible test-time settings.
4 MetaBFM
We can now combine the pieces from Section 3 into a unified update,MetaBFM, that captures the
in-distribution performance of supervised meta-RL and the out-of-distribution generalization of
unsupervised BFMs. At each training step, we first update the Latent Context Learning module with
its self-contained objective (Algorithm 1). LCL’s primary role is to augment each state with a context
estimate ˜cd, allowing the BFM networks to learn from data generated by multiple environment
dynamics (Section 3.1). LCL also produces per-timestep representations that can be passed to
off-the-shelf episodic curiosity methods to compute intrinsic rewardsri for the trajectory. We then
update the BFM networks, excluding the actorπ, over a distributionp(z). Any actions fromπ
8

Under review for the Reinforcement Learning Conference (RLC)
required for the BFM update are sampled fromπ(·|¯τ,z,α = 0,β = 0), which recovers the original
BFM policy; the remaining details are unchanged, and many ZSRL algorithms are interchangeable
here, though we use TD-JEPA. Next, we update the critic headQR with a standard off-policy
backup on weighted intrinsic and extrinsic rewards (Equation 3) and compute the actor lossLαβ-BFM
(Equation 4), both over samples fromp(α,β ). We collect trajectories from the online environment by
sampling from the same distribution ofz,α, andβ.
Supervised

Meta-RL
Zero-Shot

RL
Intrinsic