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

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

Curiosity
Figure 6:A Unified Objective.The
Lαβ-BFM loss (Eq. 4) interpolates between
zero-shot RL, intrinsic curiosity, and super-
vised meta-RL.α and β can be configured
to adjust generalization behavior.
While this update is more expensive than training the
BFM component in isolation, in practice we share net-
work backbones for efficiency. For example,π and QR
share a Transformer backbone to enable long-term mem-
ory. Many such training details are copied directly from
a strong off-policy meta-RL implementation (Grigsby
et al., 2024), and Appendix C provides additional discus-
sion. To make extrinsic rewards optional, we forceβ = 0
when collecting or training on reward-free trajectories.
Early layers that take extrinsic reward inputs toQR
andπ have their outputs multiplied byβ, allowing us to
insert and automatically ignore an arbitrary placeholder
reward.
Figure 1 showed that BFMs can learn adaptive behavior
from limited training domains where supervised meta-RL
overfits, but also that meta-RL improves with diverse
datasets while BFMs do not.MetaBFM’s main goal is to
allow for test-time configuration to take advantage of both of these methods strengths while being
able to avoid their weaknesses. Figure 7 Left returns to this experiment and shows thatMetaBFM
captures both modes: a single policy can be deployed in “ZSRL mode” (α = 0,β = 0) when meta-RL
fails to generalize, or in “reward-following mode” (α = 1,β = 1) when the dataset is diverse enough
to support it.
1 2 3 4 5 6
Adaptation Episode
0.0
0.2
0.4
0.6
0.8Oracle Normalized Return
Train T asks: 10k Envs × 10 Rewards
1 2 3 4 5 6
Adaptation Episode
Train T asks: 10k Envs × 10k Rewards
ColorGrid
Supervised Meta-RL
MetaBFM ( = 0, = 0)
MetaBFM ( : 1 0 @ Ep 3, = 0)
MetaBFM ( = 1, = 1)
0 4 8 12 16 20 24 28 32
Adaptation Episode
0.0
0.2
0.4
0.6
0.8
1.0Average success rate
T est (In-Distribution)
0 4 8 12 16 20 24 28 32
Adaptation Episode
T est (Out-of-Distribution)
Wind
Supervised Meta RL
BFM+LCL
MetaBFM ( =1 0 @ Ep 3, =0)
MetaBFM ( =0, =1)
Figure 7:Configurable MetaBFM Inference.(Left) When learning from limited training
objectives, MetaBFM in ZSRL mode (β = 0) avoids supervised meta-RL’s overfitting. (Center Left)
With more objectives, supervised mode recovers the meta-RL baseline. (Center Right) On held-out
in-distribution Wind tasks, MetaBFM remains competitive with supervised meta-RL. (Right) On
out-of-distribution Wind tasks, ZSRL-style inference adapts over long adaptation horizons.
An extension of learning more general behavior from fewer train-time objectives is generalizing to
test-time objectives that are clearly out of distribution from the supervised signals seen during
training. We study this with a variant of the Wind environment (Dorfman et al., 2020; Ni et al.,
2022), a common benchmark for OOD generalization. Agents navigate from the origin toward a
randomly placed goal,cr, while adjusting for wind in a random direction,cd. The training set
contains all wind directions but excludes goals in the upper-right quadrant. At test time, we therefore
9

Under review for the Reinforcement Learning Conference (RLC)
evaluate both held-out in-distribution tasks (any wind direction, goals in the three training quadrants)
and out-of-distribution tasks (any wind direction, goals in the excluded quadrant). Figure 7 Right
shows the adaptation curves for each set.MetaBFM significantly outperforms the supervised baseline
on the OOD test set when configured to maximize successor-feature values, while nearly matching
its performance on in-distribution tasks where train-time rewards are useful. Perhaps the most
interesting OOD result is the gap between the train-time trajectory length ofk = 2episodes and the
length at which online task inference saturates (k >20). This kind of length extrapolation is rare in
black-box meta-RL and may be enabled by offloading the usual context estimation ofcr outside of
the policy weights and into ZSRL’sz∗.
0 2 4 6 8
Adaptation Episode
0
25
50
75
100
Success Rate (%)
basketball
0 2 4 6 8
Adaptation Episode
door_open
0 2 4 6 8
Adaptation Episode
drawer_open
0 2 4 6 8
Adaptation Episode
shelf_place
0 2 4 6 8
Adaptation Episode
soccer
0 2 4 6 8
Adaptation Episode
sweep_into_goal
: 1 0 @ Ep 4,  =0
 =1,  =1
 =0,  =0
Figure 8:Online Learning in MetaWorld ML1.MetaBFM evaluated in the three main inference
configurations on six MetaWorld ML1 problems.
0
0.25
0.50
0.75
1.0
Average Success Rate (k = 10)
ML45 Train
Behavior Dataset
Best Policy in Dataset
Offline Meta-RL
Supervised Meta-RL
MetaBFM ( = 1, = 1)
Figure 9:MetaWorld ML45.Performance
across45train tasks when replay buffers are
initialized from a mixture of policies.
Finally, we evaluateMetaBFM in the more complex
continuous-control benchmark of MetaWorld (Yu
et al., 2020), using the hardware-accelerated variant
from MTBench (Joshi et al., 2025). Despite learn-
ing to maximize a space of actor objectives at train
and test time,MetaBFM converges to near-perfect per-
formance on six ML1 tasks, which measure gener-
alization over hidden-parameter variations of a sin-
gle tabletop manipulation skill. Episodic curiosity
(α = 1,β = 0) is also strongly correlated with suc-
cess rate on many ML1 tasks. Figure 9 evaluates the
full multi-task ML45 setting. Due to computational
constraints, we seed these runs with an initial replay
buffer of demonstrations generated by randomly se-
lected checkpoints from a multi-day meta-RL training
run with exploration noise enabled. We report perfor-
mance after 1M gradient steps for baselines trained
offline on this dataset, as well as meta-RL andMetaBFM agents that continue collecting online experi-
ence for exploration and coverage.MetaBFM again recovers supervised meta-RL as an approximate
lower bound on performance. However, it remains unable to generalize, in either reward-following or
ZSRL inference modes, to MetaWorld’s notoriously disjoint test set of five held-out manipulation
tasks. Because MetaBFM can learn from reward-free trajectories, expanding ML45 with additional
environment diversity may be a path toward closing this gap.
5 Conclusion
Meta-RL’s reliance on diverse reward functions has bottlenecked its scaling, even as environments
become cheaper to generate. Behavior Foundation Models sidestep this by decoupling environment
and reward representations, but standard ZSRL has kept BFMs outside meta-RL’s online adaptation
regime. Building on OptiBFM (Rupf et al., 2025), we extend the unsupervised BFM paradigm
to generalize across both novel rewards and novel environments.MetaBFM achieves this through a
self-predictive latent context for environment identification, intrinsic curiosity that drives both online
data collection and test-time exploration, and mixed supervised-unsupervised reward training that
recovers in-distribution meta-RL performance while enabling generalization beyond the training set.
10

Under review for the Reinforcement Learning Conference (RLC)
References
Siddhant Agarwal, Caleb Chuck, Harshit Sikchi, Jiaheng Hu, Max Rudolph, Scott Niekum, Peter
Stone, and Amy Zhang. A unified framework for unsupervised reinforcement learning algorithms.
InWorkshop on Reinforcement Learning Beyond Rewards @ Reinforcement Learning Conference
2025, 2025a. URLhttps://openreview.net/forum?id=tQtO75p5HA.
Siddhant Agarwal, Harshit Sikchi, Peter Stone, and Amy Zhang. Proto successor measure: Repre-
senting the behavior space of an rl agent. InInternational Conference on Machine Learning, pp.
566–586. PMLR, 2025b.
Adrià Puigdomènech Badia, Pablo Sprechmann, Alex Vitvitskyi, Daniel Guo, Bilal Piot, Steven
Kapturowski, Olivier Tieleman, Martín Arjovsky, Alexander Pritzel, Andew Bolt, et al. Never give
up: Learning directed exploration strategies.arXiv preprint arXiv:2002.06038, 2020.
Marco Bagatella, Matteo Pirotta, Ahmed Touati, Alessandro Lazaric, and Andrea Tirinzoni. Td-
jepa: Latent-predictive representations for zero-shot reinforcement learning.arXiv preprint
arXiv:2510.00739, 2025.
Marco Bagatella, Thomas Rupf, Georg Martius, and Andreas Krause. Soft forward-backward represen-
tations for zero-shot reinforcement learning with general utilities.arXiv preprint arXiv:2602.06769,
2026.
Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon
Whiteson. A survey of meta-reinforcement learning.arXiv preprint arXiv:2301.08028, 2023.
Maksim Bobrin, Ilya Zisman, Alexander Nikulin, Vladislav Kurenkov, and Dmitry Dylov. Zero-shot
adaptation of behavioral foundation models to unseen dynamics.arXiv preprint arXiv:2505.13150,
2025.
Serena Booth, W Bradley Knox, Julie Shah, Scott Niekum, Peter Stone, and Alessandro Allievi. The
perils of trial-and-error reward design: misdesign through overfitting and invalid task specifications.
InProceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 5920–5929, 2023.
Jake Bruce, Michael D Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes,
Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, et al. Genie: Generative
interactive environments. InForty-first International Conference on Machine Learning, 2024.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network
distillation.arXiv preprint arXiv:1810.12894, 2018.
Haoxuan Che, Xuanhua He, Quande Liu, Cheng Jin, and Hao Chen. Gamegen-x: Interactive
open-world game video generation.arXiv preprint arXiv:2411.00769, 2024.
Raymond Chua, Arna Ghosh, Christos Kaplanis, Blake A Richards, and Doina Precup. Learning
successor features the simple way.Advances in Neural Information Processing Systems, 37:
49957–50030, 2024.
Tri Dao, Daniel Y Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and
memory-efficient exact attention with io-awareness.arXiv preprint arXiv:2205.14135, 2022.
Hao Liu Michael Laskin Pieter Abbeel Alessandro Lazaric Lerrel Pinto Denis Yarats, David Brand-
fonbrener. Don’t change the algorithm, change the data: Exploratory data for offline reinforcement
learning.arXiv preprint arXiv:2201.13425, 2022.
Jacopo Di Ventura, Jan Felix Kleuker, Aske Plaat, and Thomas Moerland. A unified framework for
zero-shot reinforcement learning.arXiv preprint arXiv:2510.20542, 2025.
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta learning of exploration.arXiv preprint
arXiv:2008.02598, 2020.
11

Under review for the Reinforcement Learning Conference (RLC)
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. Rl2: Fast
reinforcement learning via slow reinforcement learning.arXiv preprint arXiv:1611.02779, 2016.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of
deep networks. InInternational conference on machine learning, pp. 1126–1135. PMLR, 2017.
Yunkai Gao, Rui Zhang, Jiaming Guo, Fan Wu, Qi Yi, Shaohui Peng, Siming Lan, Ruizhi Chen,
Zidong Du, Xing Hu, et al. Context shift reduction for offline meta-reinforcement learning.Advances
in Neural Information Processing Systems, 36:80024–80043, 2023.
Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Devin, Benjamin Eysenbach,
and Sergey Levine. Learning to reach goals via iterated supervised learning.arXiv preprint
arXiv:1912.06088, 2019.
Jake Grigsby, Justin Sasek, Samyak Parajuli, Daniel Adebi, Amy Zhang, and Yuke Zhu. Amago-2:
Breaking the multi-task barrier in meta-reinforcement learning with transformers.Advances in
Neural Information Processing Systems, 37:87473–87508, 2024.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena
Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar,
et al. Bootstrap your own latent-a new approach to self-supervised learning.Advances in neural
information processing systems, 33:21271–21284, 2020.
Abhishek Gupta, Benjamin Eysenbach, Chelsea Finn, and Sergey Levine. Unsupervised meta-learning
for reinforcement learning.arXiv preprint arXiv:1806.04640, 2018.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains
through world models.arXiv preprint arXiv:2301.04104, 2023.
Félix G. Harvey, Mike Yurick, Derek Nowrouzezahrai, and Christopher Pal. Robust motion in-
betweening. 39(4), 2020.
Mikael Henaff, Roberta Raileanu, Minqi Jiang, and Tim Rocktäschel. Exploration via elliptical
episodic bonuses.Advances in Neural Information Processing Systems, 35:37631–37646, 2022.
Mikael Henaff, Minqi Jiang, and Roberta Raileanu. A study of global and episodic bonuses for
exploration in contextual mdps. InInternational Conference on Machine Learning, pp. 12972–12999.
PMLR, 2023.
Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee Whye Teh, and
Nicolas Heess. Meta reinforcement learning as task inference.arXiv preprint arXiv:1905.06424,
2019.
Allan Jabri, Kyle Hsu, Abhishek Gupta, Ben Eysenbach, Sergey Levine, and Chelsea Finn. Unsuper-
vised curricula for visual meta-reinforcement learning.Advances in Neural Information Processing
Systems, 32, 2019.
Pranaya Jajoo, Harshit Sikchi, Siddhant Agarwal, Amy Zhang, Scott Niekum, and Martha White.
Regularized latent dynamics prediction is a strong baseline for behavioral foundation models.
InThe Fourteenth International Conference on Learning Representations, 2026. URLhttps:
//openreview.net/forum?id=jdL6WB5jHZ.
Scott Jeen, Tom Bewley, and Jonathan M Cullen. Zero-shot reinforcement learning from low quality
data.Advances in Neural Information Processing Systems, 37:16894–16942, 2024.
Scott Jeen, Tom Bewley, and Jonathan M Cullen. Zero-shot reinforcement learning under partial
observability.arXiv preprint arXiv:2506.15446, 2025.
Viraj Joshi, Zifan Xu, Bo Liu, Peter Stone, and Amy Zhang. Benchmarking massively parallelized
multi-task reinforcement learning for robotics tasks.arXiv preprint arXiv:2507.23172, 2025.
12