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

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

Under review for the Reinforcement Learning Conference (RLC)
Seth Karten, Rahul Dev Appapogu, and Chi Jin. Automatic generation of high-performance rl
environments.arXiv preprint arXiv:2603.12145, 2026.
Jacek Karwowski, Oliver Hayman, Xingjian Bai, Klaus Kiendlhofer, Charlie Griffin, and Joar
Max Viktor Skalse. Goodhart’s law in reinforcement learning. InThe Twelfth International
Conference on Learning Representations, 2023.
W Bradley Knox, Alessandro Allievi, Holger Banzhaf, Felix Schmitt, and Peter Stone. Reward (mis)
design for autonomous driving.Artificial Intelligence, 316:103829, 2023.
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel
Pinto, and Pieter Abbeel. URLB: Unsupervised reinforcement learning benchmark. InThirty-fifth
Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2),
2021. URLhttps://openreview.net/forum?id=lwrPkQP_is.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial,
review, and perspectives on open problems.arXiv preprint arXiv:2005.01643, 2020.
Lanqing Li, Hai Zhang, Xinyu Zhang, Shatong Zhu, Yang Yu, Junqiao Zhao, and Pheng-Ann Heng.
Towards an information theoretic framework of context-based offline meta-reinforcement learning.
Advances in Neural Information Processing Systems, 37:75642–75667, 2024.
Yitang Li, Zhengyi Luo, Tonghe Zhang, Cunxi Dai, Anssi Kanervisto, Andrea Tirinzoni, Haoyang
Weng, Kris Kitani, Mateusz Guzek, Ahmed Touati, et al. Bfm-zero: A promptable behavioral
foundation model for humanoid control using unsupervised reinforcement learning.arXiv preprint
arXiv:2511.04131, 2025.
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.02971, 2015.
Evan Z Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn. Decoupling exploration and
exploitation for meta-reinforcement learning without sacrifices. InInternational conference on
machine learning, pp. 6925–6935. PMLR, 2021.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and
Feryal Behbahani. Structured state space models for in-context reinforcement learning.arXiv
preprint arXiv:2303.03982, 2023.
Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin,
David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, et al. Isaac gym: High performance
gpu-based physics simulation for robot learning.arXiv preprint arXiv:2108.10470, 2021.
Reginald McLean, Evangelos Chatzaroulas, Luc McCutcheon, Frank Röder, Tianhe Yu, Zhanpeng
He, KR Zentner, Ryan Julian, Jordan K Terry, Isaac Woungang, et al. Meta-world+: An improved,
standardized, rl benchmark.arXiv preprint arXiv:2505.11289, 2025.
Mayank Mittal, Pascal Roth, James Tigue, Antoine Richard, Octi Zhang, Peter Du, Antonio Serrano-
Munoz, Xinjie Yao, René Zurbrügg, Nikita Rudin, et al. Isaac lab: A gpu-accelerated simulation
framework for multi-modal robot learning.arXiv preprint arXiv:2511.04831, 2025.
Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free rl can be a strong
baseline for many pomdps, 2022.
Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, Artem Agarkov, Viacheslav Sinii, and Sergey
Kolesnikov. Xland-minigrid: Scalable meta-reinforcement learning environments in jax.Advances
in Neural Information Processing Systems, 37:43809–43835, 2024.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive
coding.arXiv preprint arXiv:1807.03748, 2018.
13

Under review for the Reinforcement Learning Conference (RLC)
Octavio Pappalardo. Unsupervised learning of efficient exploration: Pre-training adaptive policies
via self-imposed goals.arXiv preprint arXiv:2601.19810, 2026.
Seohong Park, Tobias Kreiman, and Sergey Levine. Foundation policies with hilbert representations.
arXiv preprint arXiv:2402.15567, 2024.
Vitchyr H Pong, Ashvin V Nair, Laura M Smith, Catherine Huang, and Sergey Levine. Offline
meta-reinforcement learning with online self-supervision. InInternational Conference on Machine
Learning, pp. 17811–17829. Pmlr, 2022.
Martin L Puterman.Markov decision processes: discrete stochastic dynamic programming. John
Wiley & Sons, 2014.
Roberta Raileanu and Tim Rocktäschel. Ride: Rewarding impact-driven exploration for procedurally-
generated environments. InInternational Conference on Learning Representations, 2020. URL
https://openreview.net/forum?id=rkg-TJBFPB.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy
meta-reinforcement learning via probabilistic context variables. InInternational conference on
machine learning, pp. 5331–5340. PMLR, 2019.
Thomas Rupf, Marco Bagatella, Marin Vlastelica, and Andreas Krause. Optimistic task inference for
behavior foundation models.arXiv preprint arXiv:2510.20264, 2025.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach-
man. Data-efficient reinforcement learning with self-predictive representations.arXiv preprint
arXiv:2007.05929, 2020.
Max Schwarzer, Johan Samir Obando Ceron, Aaron Courville, Marc G Bellemare, Rishabh Agarwal,
and Pablo Samuel Castro. Bigger, better, faster: Human-level atari with human-level efficiency. In
International Conference on Machine Learning, pp. 30365–30380. PMLR, 2023.
Younggyo Seo, Lili Chen, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. State entropy
maximization with random encoders for efficient exploration. InInternational conference on
machine learning, pp. 9443–9454. PMLR, 2021.
Bradly C Stadie, Ge Yang, Rein Houthooft, Xi Chen, Yan Duan, Yuhuai Wu, Pieter Abbeel, and
Ilya Sutskever. Some considerations on learning to explore via meta-reinforcement learning.arXiv
preprint arXiv:1803.01118, 2018.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced
transformer with rotary position embedding.Neurocomputing, 568:127063, 2024.
Jingbo Sun, Songjun Tu, qichao Zhang, Haoran Li, Xin Liu, Yaran Chen, Ke Chen, and Dongbin Zhao.
Unsupervised zero-shot reinforcement learning via dual-value forward-backward representation.
InThe Thirteenth International Conference on Learning Representations, 2025. URL https:
//openreview.net/forum?id=0QnKnt411O.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden,
Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite.arXiv preprint
arXiv:1801.00690, 2018.
Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar
Bhoopchand, Nathalie Bradley-Schmieg, Michael Chang, Natalie Clay, Adrian Collister, et al.
Human-timescale adaptation in an open-ended task space.arXiv preprint arXiv:2301.07608, 2023.
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. InThe Thirteenth International Conference on Learning Representations.
14

Under review for the Reinforcement Learning Conference (RLC)
Ahmed Touati and Yann Ollivier. Learning one representation to 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.
Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos,
Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn.arXiv
preprint arXiv:1611.05763, 2016.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey
Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning.
InConference on robot learning, pp. 1094–1100. PMLR, 2020.
Jin Zhang, Jianhao Wang, Hao Hu, Tong Chen, Yingfeng Chen, Changjie Fan, and Chongjie Zhang.
Metacure: Meta reinforcement learning with empowerment-driven exploration. InInternational
Conference on Machine Learning, pp. 12600–12610. PMLR, 2021.
Chongyi Zheng, Royina Karegoudra Jayanth, and Benjamin Eysenbach. Can we really learn one
representation to optimize all rewards?arXiv preprint arXiv:2602.11399, 2026.
Ce Zhou, Qian Li, Chen Li, Jun Yu, Yixin Liu, Guangjing Wang, Kai Zhang, Cheng Ji, Qiben Yan,
Lifang He, et al. A comprehensive survey on pretrained foundation models: A history from bert to
chatgpt.International Journal of Machine Learning and Cybernetics, 16(12):9851–9915, 2025.
Luisa Zintgraf.Fast adaptation via meta reinforcement learning. PhD thesis, University of Oxford,
2022.
Luisa Zintgraf, Sebastian Schulze, Cong Lu, Leo Feng, Maximilian Igl, Kyriacos Shiarlis, Yarin
Gal, Katja Hofmann, and Shimon Whiteson. Varibad: Variational bayes-adaptive deep rl via
meta-learning.Journal of Machine Learning Research, 22(289):1–39, 2021a.
Luisa M Zintgraf, Leo Feng, Cong Lu, Maximilian Igl, Kristian Hartikainen, Katja Hofmann, and
Shimon Whiteson. Exploration in approximate hyper-state space for meta reinforcement learning.
InInternational Conference on Machine Learning, pp. 12991–13001. PMLR, 2021b.
Appendix
A Additional Background
A.1 Reinforcement Learning
A Markov Decision Process (MDP) (Puterman, 2014) is defined as a tuple(S,A,µ,P,r )where S is
called the state space,A is the action space,µ is the initial state distribution,P :S×A→ ∆(S)
is the transition function that maps a given state and action to a distribution of next states and
r :S→R is the scalar reward function. The agent executes a policy,π :S→ ∆(A), to interact
with the environment. The Reinforcement Learning (RL) objective is to find the optimal policy,π∗,
defined as,
π∗ =argmax πEπ
[∑
t
γtr(st)
]
(5)
15

Under review for the Reinforcement Learning Conference (RLC)
A.2 Behavior Foundation Models
A Behavior Foundation Model(BFM), for a given dynamical system, allows an agent to produce
near-optimal policies for any given reward function using little or no additional computation. As a
result, this paradigm has often been referred to asZero-Shot Reinforcement Learning. Specifically, in
this work, we consider successor-measure-based BFMs that learn to represent successor measures for
a wide range of policies. Mathematically, the successor measure for a policy is defined as,
Mπ(s,a,X) =E π
[∞∑
t=0
γtPπ(st∈X|s 0 =s,a 0 =a)
]
,∀X⊂S.(6)
BFMs represent Mπz (successor measure for policy parameterized usingz) as Mπz(s,a,s +) =
ψ(s,a,z )Tϕ(s+). Connections have been drawn (Touati & Ollivier, 2021; Agarwal et al., 2025b;a)
between this representation and successor features withψ being the successor feature for the state
feature Eρ[ϕϕT ]ϕ(s). Methods like (Touati & Ollivier, 2021; Agarwal et al., 2025b; Jajoo et al., 2026)
use the following objective to learn these representations,
L(ϕ,ψ) =−E s,a,s′∼ρ[ψ(s,a,z) Tϕ(s′)]
+ 1
2 Es,a,s′∼ρ,s+∼ρ[(ψ(s,a,z) Tϕ(s+)−γ ¯ψ(s′,πz(s′),z)T ¯ϕ(s+))2](7)
where ρ is the offline dataset. These BFMs consider tasks defined by rewards,rz(s) =ϕ(s)Tz for
which the corresponding optimal Q function can be written as,Q∗
z(s,a ) = ψ(s,a,z )Tz with the
optimal policyπz = argmaxaQ∗
z(s,a ) = argmaxaψ(s,a,z )Tz. Overall, BFMs can be represented
using( ϕ,ψ,π z). All the three networksϕ, ψ and πz are learnt using reward-free interactions withϕ
andψbeing trained using Equation 7 andπz as the optimal policy for the rewardrz(s) =ϕ(s) Tz.
At inference, when a reward function is provided, the BFMs simply need to find thez for the given
reward function as the corresponding policyπz is already trained to be optimal forrz. Thisz can
be inferred simply by linear regression on reward samples, i.e.z∗ = argminzEρ[(r(s)−ϕ (s)Tz)2].
Several methods (Touati & Ollivier, 2021; Agarwal et al., 2025b; Jajoo et al., 2026; Park et al., 2024;
Bagatella et al., 2025) have been proposed with a similar construction of BFMs.
B Environment Details
B.1 Toy Meta-RL
B.1.1 ColorGrid
ColorGrid is ann×n continuous-control gridworld in which the agent acts on grid of colors, with a
hidden reward gridcr defining a scalar payoff at each cell. Rather than observing this underlying
color grid, the agent’s observations are produced from a permuted color grid produced by a hidden
cell permutationcd. Each task is therefore a pair(cd,cr), a construction that disentangles two
independent sources of partial observability the agent must reason about within a trial.
Dynamics and reward.The agent’s true state is a continuous positionst∈ [0,n ]2, advanced
by a 2D velocity actionat ∈ [−1, 1]2 with st+1 = clip(st +at, 0,n ), and reward is read from
the true grid as rt = R[⌊s(y)
t ⌋,⌊s (x)
t ⌋]. Actions therefore act faithfully on the true grid, and
only the position readout to the agent is scrambled bycd. The integer part of the agent’s po-
sition is relabelled according tocd, while the sub-cell fraction passes through unchanged. Smooth
motion within a true cell therefore looks smooth in observation space, but the moment the
16