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

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

Under review for the Reinforcement Learning Conference (RLC)
Meta-Reinforcement Learning with Zero-Shot Rein-
forcement Learning
Anonymous authors
Paper under double-blind review
Abstract
Meta-reinforcement learning (meta-RL) agents adapt to novel tasks from test-time
experience, but require diverse training sets of environments and reward functions
that are expensive to construct. Behavior Foundation Models (BFMs) learn policies
from reward-free data that are capable of zero-shot RL (ZSRL): adapting to new
objectives at test time when given a large reward-labeled dataset. Recent work has
shown that BFM task inference can be performed online without prior data, making
BFMs and meta-RL direct competitors for generalization in fixed environments. We
ask whether BFMs can be extended to adapt to novel reward functions in novel
environments, and identify four key limitations: environment identification, online
data collection, test-time exploration-exploitation, and mixed reward supervision.
We study these subproblems through toy domains, standard BFM benchmarks, and
simulated humanoid locomotion, then propose a unified framework that addresses
all four. The resulting method,MetaBFM, is a hybrid RL agent spanning meta-RL,
ZSRL, and intrinsic exploration. During training,MetaBFM combines supervised
reward-following with unsupervised reward-free learning; at test time, it interpolates
between exploration, exploitation, and ZSRL-style reward inference. We evaluate
MetaBFM on two toy meta-RL domains and at scale in MetaWorld, showing that
hybrid meta-RL/ZSRL agents can learn more general behavior from the same
set of reward functions and may reduce the need for future meta-RL domains to
hand-design diverse training sets.
1 Introduction
Meta-Reinforcement Learning (RL) agents improve on novel tasks by adapting their decision-making
from test-time experience (Beck et al., 2023). This capability relies on diverse training sets of
environments and reward functions, and benefits from online data collection over many trials. Despite
much algorithmic progress, meta-RL remains limited by the cost of generating these training sets.
Most work focuses on toy problems where similar environments and/or reward functions are easily
sampled; large-scale applications like AdA (Team et al., 2023), where a team of developers repurposed
a commercial game engine to generate complex tasks and curricula, are rare exceptions. Generating a
diverse set of learnablereward functionsthat covers the desired range of behavior may be inherently
difficult. Reward functions are challenging to design at scale (Booth et al., 2023; Knox et al., 2023),
and imperfect proxies can break down under optimization pressure (Karwowski et al., 2023). However,
with hardware-accelerated simulators (Makoviychuk et al., 2021; Mittal et al., 2025), LLM-assisted
coding (Karten et al., 2026), and promptable world models (Bruce et al., 2024; Che et al., 2024), the
cost of generatingenvironmentsis rapidly declining.
Behavior Foundation Models (BFMs) are an emerging class of unsupervised RL methods (Laskin
et al., 2021) that train in a reward-free environment by learning to maximize many possible objectives
(Agarwal et al., 2025a). At test time, a target objective is specified by model inference on a large
dataset of trajectories labeled with a particular reward function. BFMs’ unsupervised training
and supervised testing datasets are loosely analogous to the pretraining and finetuning paradigm
1

Under review for the Reinforcement Learning Conference (RLC)
of foundation models in other areas (Zhou et al., 2025). This adaptation process is often called
zero-shot RL (ZSRL) (Touati et al., 2022) and has shown promising results in domains like humanoid
locomotion (Tirinzoni et al.; Li et al., 2025). Meta-RL and BFMs both enable generalization without
additional training, but ZSRL requires a large supervised adaptation dataset in advance, while
meta-RL assumes the policy actively collects its own (far smaller) dataset. This distinction makes the
two approaches applicable to different types of adaptation problems, and their test-time performance
is difficult to compare directly. However, recent work by Rupf et al. (2025) demonstrated that
ZSRL task inference can be done online, without a preexisting dataset, under the same conditions as
meta-RL. As a result, meta-RL and BFMs can now be considered direct competitors on problems
that require adaptation to novel reward functionsin a fixed environment.
If unsupervised BFMs could behave like supervised meta-RL in the broader setting of generalizing to
novel reward functionsin novel environments, their ability to avoid meta-RL’s demand for train-time
reward functions could help extend generalist RL to more diverse domains. We argue that this
requires BFMs to (1) identify and generalize to novel environments, (2) benefit from large-scale
online data collection, (3) make exploration-exploitation tradeoffs at test time, and (4) improve when
we relax the assumption that no reward functions are available during training. We study each of
these subproblems in turn with detailed experiments on a challenging toy problem, standard BFM
benchmarks, and large-scale simulated humanoid locomotion, arriving at a unified framework that
addresses all four. The resulting method,MetaBFM, trains a hybrid RL agent that spans meta-RL,
ZSRL, and intrinsic exploration. During training,MetaBFM combines supervised reward-following
with unsupervised reward-free learning. At test time, the policy can interpolate between exploration,
exploitation, and ZSRL-style reward inference. The main benefit is thatMetaBFM retains meta-RL’s
ability to adapt to in-distribution environments and objectives, while adding the possibility of
generalizing to out-of-distribution reward functions. In practice, this allows the agent to learn more
from a fixed set of supervised rewards and shifts the burden from designing large training sets of
reward functions to designing diverse environments.
2 Background
2.1 Meta-Reinforcement Learning
A meta-RL policyπ can be evaluated by its cumulative return over the firstk attempts, or episodes,
in a new task. This objective creates an exploration-exploitation tradeoff at test time: the policy may
sacrifice performance in early attempts to learn more about the task if exploiting that knowledge
later leads to higher cumulative return (Zintgraf, 2022). Methods differ in whether this tradeoff
is managed by the designer as a hyperparameter (Finn et al., 2017) or learned by the policy itself
(Duan et al., 2016). Policies may also be givenk free episodes to explore without consequence before
being evaluated on attemptk + 1; thisk-shotsetting creates explicit exploration and exploitation
phases (Liu et al., 2021; Stadie et al., 2018).
Meta-RL’s test-time adaptation is enabled by training on a diverse set of tasks. In this paper, a
“task” refers specifically to a combination of an environment and an objective—or, equivalently, to
a pair of parameters(cd,cr)that alter the transition dynamics and reward function, respectively.
The basic goal of a context-based meta-learner is to use the multi-episodic trajectory1 of interaction
with the current task so far,τ:t ={s0,a 0,r 1,s 1,a 1,...,s t}, to infer the unknown(cd,cr), exploit the
unique characteristics of the task, and maximize returns. Methods broadly differ in whether the
parameters( cd,cr)are estimated by supervised regression to ground-truth values (Humplik et al.,
2019), as the byproduct of a self-supervised objective such as dynamics modeling (Zintgraf et al.,
2021a), or implicitly through the objective of maximizing returns (Wang et al., 2016). In each case,
the idea is to establish a training objective that is minimized by correctly distinguishing(cd,cr)and
then conditioningπon the resulting representation.
1The main text assumes common MDP RL notation. Appendix A.1 provides a reference and a standard introduction.
2

Under review for the Reinforcement Learning Conference (RLC)
Meta-RL is at its best when we can afford sample-inefficient online data collection (Team et al., 2023;
Nikulin et al., 2024; Lu et al., 2023), which may be acceptable when we are willing to trade train-time
compute for efficient test-time adaptation. Meta-learning is harder when restricted to a fixed dataset
(Pong et al., 2022; Dorfman et al., 2020; Gao et al., 2023; Li et al., 2024), as it suffers from an
extreme form of the distribution shift problem that defines offline RL more broadly (Levine et al.,
2020). Methods also differ in whether task rewards are available during meta-training. Standard
meta-RL issupervisedby a distribution of reward-labeled tasks, whereasunsupervisedmeta-RL
learns from reward-free interaction and only encounters rewards at test time (Beck et al., 2023).
We focus on online meta-RL that spans both regimes, meta-training on a set of reward functions
while simultaneously learning a reward-free BFM that supports ZSRL-style task inference at test
time. This hybrid framing leads to a method distinct from prior unsupervised meta-RL (Gupta et al.,
2018; Jabri et al., 2019; Pappalardo, 2026), which omits reward supervision entirely, and from prior
ZSRL/BFM work, which does not meta-train across environments.
2.2 Behavior Foundation Models
BFMs are unsupervised RL agents based on successor features (Agarwal et al., 2025a; Di Ventura
et al., 2025). The idea is to learn a latent-conditioned policyπ(a|s,z )from a reward-free dataset and
infer thez that maximizes a particular task of interest from a reward-labeled dataset, without further
training (Jeen et al., 2025). In general, we letϕ :S→R d define latent rewardsrz(s) = ϕ(s)⊤z
and learn successor featuresψ(s,a,z )≈E πz[∑
t≥0γtϕ(st+1)|s 0 =s,a 0 =a], so thatQSF(s,a,z ) =
ψ(s,a,z )⊤z can serve as the critic in a standard off-policy actor-critic update (Lillicrap et al., 2015)
over some distribution ofz, e.g.,Lactor =−Es∼D,z∼p(z),a∼π(·|s,z) [QSF(s,a,z)] . In this work, the
details of trainingψ and ϕ will be interchangeable and are not necessary to follow the main text,
so we defer a more complete introduction to Appendix A.2. Most recent work is based on the
FB algorithm (Touati & Ollivier, 2021; Touati et al., 2022), but there are many alternatives and
extensions (Chua et al., 2024; Bagatella et al., 2026; Zheng et al., 2026; Jeen et al., 2024; Agarwal
et al., 2025b; Park et al., 2024). Our experiments primarily use TD-JEPA (Bagatella et al., 2025)
because it shares similarities with off-policy supervised RL and lets us transfer implementation details
from an optimized meta-RL baseline when appropriate.
BFMs benefit from diverse training datasets with broad state-action coverage, and are most commonly
applied to offline datasets originally collected by a policy that optimized for exploration (Burda et al.,
2018; Laskin et al., 2021). They adapt to a test-time task by fittingz∗ = arg minz
∑
Dr(ri−ϕ(si)⊤z)2
on a reward-labeled dataset to deployπ(·|s,z ∗), and we will refer to this inference process as
zero-shot RL (ZSRL). Meta-RL’s inference ofcr plays a similar role as ZSRL’s inference ofz∗: both
parameters estimate the active reward function and therefore enable generalization to novel objectives.
ZSRL relies on an existingofflinedatasetDr to fitz∗, after which this parameter is fixed. In contrast,
meta-RL typically inferscr fromτ:t (anonline Dr it has collected itself) by updating a belief between
timesteps or episodes (Rakelly et al., 2019). This discrepancy makes ZSRL and meta-RL difficult
to compare directly. In fact, “zero-shot RL” is, by design, unsuited to zero-shot generalization as
typically defined (Section 2.1), where we expect a policy to perform well in a novel task on the first
attempt with no prior data. OptiBFM (Rupf et al., 2025) helps to close this gap by demonstrating
that, as in meta-RL, we can maintain a belief overz∗ that is updated between timesteps or episodes.
2.3 Generalization of BFMs and Supervised Meta-RL
With OptiBFM’sonline task inferenceenabled, BFMs and meta-RL can be directly compared. We
will use a didactic ColorGrid problem as a running example through Section 3. ColorGrid creates an
n×n continuous 2D gridworld where the agent receives a reward based on the unit square cell of
its current location, and the objective is to navigate to the cell with the highest reward and remain
there. The underlying cell layout can be randomly permuted such that any two cells may be adjacent.
In short, we can generate(n2)!distinct environments where nearly every transition is unpredictable
3

Under review for the Reinforcement Learning Conference (RLC)
0.00
0.25
0.50
0.75
1.00
Norm. Return
1 10 100 10,000
0.00
0.25
0.50
0.75
1.00
Norm. Return
1 2 3 4 5 6
Adaptation Episode
0.00
0.25
0.50
0.75
1.00
Norm. Return
1 2 3 4 5 6
Adaptation Episode
1 2 3 4 5 6
Adaptation Episode
1 2 3 4 5 6
Adaptation Episode
Meta-RL w/o Seq Model
Meta-RL w/o Reward Input
1010010,000
Environments Available During Training
Rewards Available During Training
 ZSRL does not utilize training rewards
Train T est (Held-Out Envs and Rewards)
Supervised Meta-RL ZSRL w/ Privileged T ask Inference ZSRL w/ Online T ask Inference
Train T est (Held-Out Envs and Rewards)
Supervised Meta-RL ZSRL w/ Privileged T ask Inference ZSRL w/ Online T ask Inference
Figure 1:BFM and Meta-RL Generalization.In offline ColorGrid (n = 9), we vary training
environments (columns) and rewards (rows), then evaluatek = 6episode adaptation on train and
held-out tasks. Meta-RL improves with training set diversity but needs many reward functions to
generalize; BFMs adapt in the fixed-environment setting (left column), but fail across environments.
without a reasonable estimate of the active grid permutation,cd. Randomly sampling a reward for
each cell generates infinitely manycr values. Additional environment details are in Appendix B.
As a first experiment, we splitcd and cr into train/test sets and generate tasks from their cross
product, measuring environment and reward generalization separately. To replicate BFMs’ high-
coverage offline setting, we collect large random-policy datasets. Our BFM baseline is the better of
FB and TD-JEPA, while our meta-RL baseline is an off-policy variant of RL2 (Duan et al., 2016)
that makes offline training reasonable, if still slightly unfair—we return to this later. Figure 1 reports
results forn = 9. Meta-RL follows a predictable trend: adding environments or reward functions
improves generalization, but closing the train/test gap requires a diverse mixture of both. In this
sense, supervised meta-RL is designed for the bottom-right corner of Fig. 1. OptiBFM allows BFMs
to match the standardz∗ inference procedure, though adaptation barely saturates within thek = 6
test episodes. In contrast, meta-RL has trained on reward functions and is incentivized to learn a
narrow prior overcr. For now, we are not concerned with the performance gap between BFMs and
meta-RL where BFMs show meaningful adaptation. Instead, we focus on BFMs’ relative performance:
the train curves in the left column of Fig. 1 represent the established BFM setting, and performance
holds as evaluated reward functions grow in a fixed environment, even though the reward labels
go unused and therefore act as a test set themselves. However, BFMs immediately collapse with
multiple training environments, a problem also identified by Bobrin et al. (2025).
3 BFMs as Meta-RL Agents
BFMs can adapt to unseen reward functions overk attempts without train-time rewards. However, to
be useful in meta-RL, they must also generalize across environments. This is only one of several gaps:
meta-RL benefits from online training, while BFMs default to offline datasets; meta-RL manages
exploration and exploitation overk test episodes, while BFMs do not; and complete reward-free
learning is often unnecessary in simulated online RL domains. We can usually define some reward-
4