# Towards an Information Theoretic Framework of Context-Based Offline Meta-
Reinforcement Learning

Lanqing Li1,2 , Hai Zhang3∗, Xinyu Zhang4, Shatong Zhu3,
Junqiao Zhao3 , Pheng-Ann Heng1
1The Chinese University of Hong Kong 2Zhejiang Lab
3Tongji University 4State University of New York at Stony Brook
<EMAIL_ADDRESS><EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS>Equal Contribution. Work
done during Hai Zhang’s internship at Zhejiang Lab.Corresponding author

###### Abstract

As a marriage between offline RL and meta-RL, the advent of offline meta-
reinforcement learning (OMRL) has shown great promise in enabling RL agents to
multi-task and quickly adapt while acquiring knowledge safely. Among which,
Context-based OMRL (COMRL) as a popular paradigm, aims to learn a universal
policy conditioned on effective task representations. In this work, by
examining several key milestones in the field of COMRL, we propose to
integrate these seemingly independent methodologies into a unified information
theoretic framework. Most importantly, we show that the pre-existing COMRL
algorithms are essentially optimizing the same mutual information objective
between the task variable $\bm{M}$ and its latent representation $\bm{Z}$ by
implementing various approximate bounds. Based on the theoretical insight and
the information bottleneck principle, we arrive at a novel algorithm dubbed
UNICORN, which exhibits remarkable generalization across a broad spectrum of
RL benchmarks, context shift scenarios, data qualities and deep learning
architectures, attaining the new state-of-the-art. We believe that our
framework could open up avenues for new optimality bounds and COMRL
algorithms.

##  1 Introduction

The ability to swiftly learn and generalize to new tasks is a hallmark of
human intelligence. In pursuit of this high level of artificial intelligence
(AI), the paradigm of meta-reinforcement learning (RL) proposes to train AI
agents in a trial-and-error manner by interacting with multiple external
environments. In order to quickly adapt to the unknown, the agents need to
integrate prior knowledge with minimal experience (namely the context)
collected from the new tasks or environments, without over-fitting to the new
data. This meta-RL mechanism has been adopted in many applications such as
games (Zhao et al., 2022a; Beck et al., 2023), robotics (Yu et al., 2020a;
Kumar et al., 2021) and drug discovery (Feng et al., 2022).

However, for data collection, classical meta-RL usually requires enormous
online explorations of the environments (Finn et al., 2017; Rakelly et al.,
2019), which is impractical in many safety-critical scenarios such as
healthcare (Gottesman et al., 2019) and autonomous driving (Shalev-Shwartz et
al., 2016). As a remedy, offline RL (Levine et al., 2020) enables agents to
learn from logged experience only, therefore circumventing risky or costly
online interactions. Recently, offline meta-RL (OMRL) (Li et al., 2021b;
Mitchell et al., 2021; Dorfman et al., 2021) has emerged as a novel paradigm
to significantly extend the applicable range of RL by "killing two birds in
one stone": it builds powerful agents that can quickly learn and adapt by
meta-learning, while leveraging offline RL mechanism to ensure a safe and
efficient optimization procedure. In the context of supervised or self-
supervised learning, which is de facto offline, OMRL is reminiscent of the
multi-task learning (Raffel et al., 2020), meta-training (Finn et al., 2017)
and fine-tuning (Hu et al., 2021; Li and Liang, 2021; Lester et al., 2021) of
pre-trained large models. We envision it to be a cornerstone of RL foundation
models (Bommasani et al., 2021; Reed et al., 2022) in the future.

Along the line of OMRL research, context-based offline meta-reinforcement
learning (COMRL) is a popular paradigm that seeks optimal meta-policy
conditioning on the context of Markov Decision Processes (MDPs). Intuitively,
the crux of COMRL is learning effective task representations, and therefore
enabling the agent to react optimally and adaptively in various contexts. To
achieve this goal, one of the earliest COMRL algorithms FOCAL (Li et al.,
2021b) proposes to capture the structure of task representations by distance
metric learning. From a geometric perspective, it essentially performs
clustering by pushing away latent embeddings of different tasks while pulling
together those from the same task, hence ensuring consistent and
distinguishable task representations.

Despite its effectiveness, FOCAL is reported to be vulnerable to context
shifts (Li et al., 2021a), i.e., when testing on out-of-distribution (OOD)
data. Such problems are particularly challenging for OMRL, since any context
shift incurred at test time can not be rectified in the fully offline setting,
which may result in severely degraded generalization performance (Li et al.,
2021a; Gao et al., 2023). To alleviate the problem, follow-up works such as
CORRO (Yuan and Lu, 2022) reformulates the task representation learning of
COMRL as maximizing the mutual information $I(\bm{Z};\bm{M})$ between the task
variable $\bm{M}$ and its latent $\bm{Z}$. It then approximates
$I(\bm{Z};\bm{M})$ by an InfoNCE (Oord et al., 2018) contrastive loss, where
the positive and negative pairs are conditioned on identical state-action
tuples $(\bm{s},\bm{a})$. Inspired by CORRO, a recently proposed method CSRO
(Gao et al., 2023) introduces an additional mutual information term between
$\bm{Z}$ and $(\bm{s},\bm{a})$. By explicitly minimizing it along with the
FOCAL objective, CSRO is demonstrated to achieve the state-of-the-art (SOTA)
generalization performance on various MuJoCo (Todorov et al., 2012)
benchmarks.

In this paper, following the recent development and storyline of COMRL, we
present a novel and unified view of several signature methods including FOCAL,
CORRO and CSRO under a information theoretic framework. We first show that
these algorithms are essentially optimizing different approximate bounds of
the same mutual information objective $I(\bm{Z},\bm{M})$. Then based on our
proposed Unified Information Theoretic Framework of Context-Based Offline
Meta-Reinforcement Learning (UNICORN), we arrive at a new COMRL optimization
scheme by deriving a unified and generalized task representation learning
objective in light of the information bottleneck principle (Tishby et al.,
1999; Tishby and Zaslavsky, 2015). We demonstrate experimentally that the
proposed method achieves state-of-the-art in-distribution and OOD
generalization performance on a wide range of RL domains, OOD settings, data
qualities and model architectures. Our framework potentially provides a
principled roadmap to new COMRL algorithms by seeking better optimality
bounds/approximations of the objective, as well as new implementations to
further combat context shift.

##  2 Related Works

Offline Meta-Reinforcement Learning. OMRL solves the problems of
generalization and expensive data collection by meta-learning on pre-collected
offline datasets. To overcome the overestimation of out-of-distribution
actions (Kumar et al., 2019; Levine et al., 2020; Gao et al., 2022), offline
RL usually leverages regularizations such as KL-divergence (Jaques et al.,
2019), maximum mean discrepancy (MMD) (Kumar et al., 2019) and Wasserstein
Distance (Gulrajani et al., 2017) to constrain the distribution mismatch
between the behavior and learner policies. Moreover, a general mechanism
called behavior regularized actor critic (BRAC) (Wu et al., 2019) involving
value penalty and behavior regularization has proven successful in several
previous studies (Li et al., 2021b; Gao et al., 2023), and is adopted in our
method.

From a meta-learning perspective, the realization of OMRL can be categorized
as gradient-based and context-based paradigms. Gradient-based meta-RL methods
(Finn et al., 2017; Mitchell et al., 2021) embed an optimization procedure
such as gradient descent into the meta-level iteration loops, which are model-
agnostic but tend to reach lower asymptotic performance with less
computational efficiency in practice. COMRL conditions the models on a
contextual representation of MDP (Hallak et al., 2015) generated by recurrent
(Wang et al., 2016; Fakoor et al., 2020), recursive (Mishra et al., 2017; Xu
et al., 2022), probabilistic (Rakelly et al., 2019; Zintgraf et al., 2019) or
deterministic (Li et al., 2021b, a; Yuan and Lu, 2022) task inference modules.
At test time, it leverages neural networks as universal value/policy function
approximators (Schaul et al., 2015) to perform meta-adaptation by transfer
learning of the task representations. Depending on the source of the testing
data, fully offline (Li et al., 2021b, a; Yuan and Lu, 2022) and offline-to-
online (Zhao et al., 2022b; Pong et al., 2022) COMRL have been investigated.
Notwithstanding, both paradigms may suffer significant OOD as the data-
collecting policies might drastically differ from the training sets. In this
work, we restrict our attention to fully offline COMRL, and demonstrate the
remarkable robustness of our proposed framework to context shift in this
setting.

##  3 Method

Figure 1: Context shift of COMRL in Ant-Dir. Left: Given a task specified by a
goal direction (dashed line), the RL agent is trained on data generated by a
variety of behavior policies trained on the same task (red). At test time,
however, the context might be collected by behavior policies trained on
different tasks (blue), causing a context shift of OOD behavior policies
(Section 4.4). Middle: Under OOD testing context, UNICORN (red) is more robust
than baselines such as FOCAL (green) in terms of navigating the Ant robot
along the right direction. Right: Besides behavior policy, the task
distribution (e.g., goal positions in Ant) can induce significant context
shift (Section 5.2), which is also a challenging scenario for COMRL models to
generalize.

###  3.1 Preliminaries and Problem Statement

We consider MDP modeled as
$M=(\mathcal{S},\mathcal{A},T,R,\mathcal{\rho}_{0},\gamma)$ with state space
$\mathcal{S}$, action space $\mathcal{A}$, transition function
$T(\bm{s}^{\prime}|\bm{s},\bm{a})$, bounded reward function
$R(\bm{s},\bm{a})$, initial state distribution $\rho_{0}(\bm{s})$ and discount
factor $\gamma\in(0,1)$. The goal is to find a policy
$\pi:\mathcal{S}\rightarrow\mathcal{A}$ to maximize the expected return.
Starting from the initial state, for each time step, the agent performs an
action sampled from $\pi$, then the environment updates the state with $T$ and
returns a reward with $R$. We denote the marginal state distribution at time
$t$ as $\mu_{\pi}^{t}(\bm{s})$. The V-function and Q-function are given by

$\displaystyle V_{\pi}(\bm{s})$
$\displaystyle=\sum_{t=0}^{\infty}\gamma^{t}\mathbb{E}_{\bm{s}_{t}\sim\mu_{\pi}^{t}(\bm{s}),\bm{a}_{t}\sim\pi}[R(\bm{s}_{t},\bm{a}_{t})]$
(1)

$\displaystyle Q_{\pi}(\bm{s},\bm{a})$
$\displaystyle=R(\bm{s},\bm{a})+\gamma\mathbb{E}_{\bm{s}^{\prime}\sim
T(\bm{s}^{\prime}|\bm{s},\bm{a})}[V_{\pi}(\bm{s}^{\prime})].$ (2)

In this paper, we restrict to scenarios where tasks share the same state and
action space while differing in reward and transition functions, with each
OMRL task defined as an instance of MDP:
$M^{i}=(\mathcal{S},\mathcal{A},T^{i},R^{i},\rho_{0},\gamma)\in\mathcal{M}$,
where $\mathcal{M}$ is the set of all possible MDPs to be considered. For each
task index $i$, the offline dataset
$X^{i}=\\{(\bm{s}_{j}^{i},\bm{a}_{j}^{i},r_{j}^{i},\bm{s}{{}^{\prime}}_{j}^{i})\\}$
is collected by a behavior policy $\pi^{i}_{\beta}$. For meta-learning, given
a trajectory segment
$\bm{c}^{i}_{1:n}=\\{(\bm{s}_{j}^{i},\bm{a}_{j}^{i},r_{j}^{i},\bm{s}{{}^{\prime}}_{j}^{i})\\}_{j=1}^{n}$
of length $n$ as the context of $M^{i}$, COMRL employs a context encoder
$q_{\bm{\phi}}(\bm{z}|\bm{c}_{1:n}^{i})$ to obtain the latent representation
$\bm{z}^{i}$ of task $M^{i}$. If $\bm{z}$ contains enough information about
the task identity, COMRL can be treated as a special case of RL on partially-
observed MDP (Kaelbling et al., 1998), where $\bm{z}$ is interpreted as a
faithful representation of the unobserved state. By conditioning on $\bm{z}$,
the learning of universal policy $\pi_{\bm{\theta}}(\cdot|\bm{s},\bm{z})$ and
value function $V_{\pi}(\bm{s},\bm{z})$ (Schaul et al., 2015) become regular
RL, and optimality can be attained by Bellman updates with theoretical
guarantees. To this end, FOCAL (Li et al., 2021b) proposes to decouple COMRL
into the upstream task representation learning and downstream offline RL. For
the former, FOCAL effectively clusters task embeddings by metric learning:

$\displaystyle\mathcal{L}_{\text{FOCAL}}=$
$\displaystyle\,\,\underset{\bm{\phi}}{\min}\,\,\mathbb{E}_{i,j}\text{\LARGE\\{}\textbf{1}\\{i=j\\}||\bm{z}^{i}-\bm{z}^{j}||_{2}^{2}+\textbf{1}\\{i\neq
j\\}\frac{\beta}{||\bm{z}^{i}-\bm{z}^{j}||_{2}^{n}+\epsilon}\text{\LARGE\\}}.$
(3)

However, it is empirically shown that FOCAL struggles to generalize in the
presence of context shift (Li et al., 2021a). This problem has been formulated
in several earliest studies of OMRL. Li et al. (2020) observed that when there
are large divergence between the state-action distributions of the offline
datasets, due to shortcut learning (Geirhos et al., 2020), the task encoder
may learn to ignore the causal information like rewards and spuriously
correlate primarily state-action pairs to the task identity, leading to poor
generalization performance. Dorfman et al. (2021) concurrently identified a
related problem which they termed MDP ambiguity in the context of Bayesian
offline RL. An example is illustrated in Figure 1. The problem is further
exacerbated in the fully offline setting, as the testing distribution is fixed
and cannot be augmented by online exploration, allowing no theoretical
guarantee for the context shift.

To mitigate context shift, a subsequent algorithm CORRO (Yuan and Lu, 2022)
proposes to optimize a lower bound of $I(\bm{Z};\bm{M})$ in the form of an
InfoNCE (Oord et al., 2018) contrastive loss

$\displaystyle\mathcal{L}_{\text{CORRO}}$
$\displaystyle=\underset{\bm{\phi}}{\min}\,\,\mathbb{E}_{\bm{x},\bm{z}}\left[-\log\left(\frac{h(\bm{x},\bm{z})}{\sum_{M^{*}\in\mathcal{M}}h(\bm{x}^{*},\bm{z})}\right)\right],$
(4)

where $\bm{x}=(\bm{s},\bm{a},\bm{r},\bm{s}^{\prime})$, $\bm{z}\sim
q_{\bm{\phi}}(\bm{z}|\bm{x})$,
$h(\bm{x},\bm{z})=\frac{P(\bm{z}|\bm{x})}{P(\bm{z})}\approx\frac{q_{\bm{\phi}}(\bm{z}|\bm{x})}{P(\bm{z})}$,
and $\bm{x}^{*}=(\bm{s},\bm{a},r^{*},\bm{s}^{\prime*})$ as a transition tuple
generated for task $M^{*}\in\mathcal{M}$ conditioned on the same
$(\bm{s},\bm{a})$. To further combat the spurious correlation between the task
representation and behavior-policy-induced state-actions, a recently proposed
method CSRO (Gao et al., 2023) introduces a CLUB upper bound (Cheng et al.,
2020) of the mutual information between $\bm{z}$ and $(\bm{s},\bm{a})$, to
regularize the FOCAL objective:

$\displaystyle\mathcal{L}_{\text{CSRO}}$
$\displaystyle=\underset{\bm{\phi}}{\min}\left\\{\mathcal{L}_{\text{FOCAL}}+\lambda
L_{\text{CLUB}}\right\\}$ (5) $\displaystyle L_{\text{CLUB}}$
$\displaystyle=\mathbb{E}_{i}\left[\log
q_{\bm{\phi}}(\bm{z}_{i}|\bm{s}_{i},\bm{a}_{i})-\mathbb{E}_{j}\left[\log
q_{\bm{\phi}}(\bm{z}_{j}|\bm{s}_{i},\bm{a}_{i})\right]\right].$ (6)

In the next section, we will show how these algorithms are inherently
connected and can be further improved under our proposed unified information
theoretic framework.

###  3.2 A Unified Information Theoretic Framework

We start with a formal definition of task representation learning in COMRL:

######  Definition 3.1.

Given an input context variable $\bm{X}\in\mathcal{X}$ and its associated
task/MDP random variable $\bm{M}\in\mathcal{M}$, task representation learning
in COMRL aims to find a minimal sufficient statistics $\bm{Z}$ of $\bm{X}$
with respect $\bm{M}$.

In pure statistical terms, Definition 3.1 implies that an ideal representation
$\bm{Z}$ is the simplest mapping of $\bm{X}$ that captures the mutual
information $I(\bm{X};\bm{M})$. We therefore assume the following dependency
structures in terms of directed graphical models:

######  Assumption 3.2.

The dependency graphs of COMRL are given by

satisfying $I(\bm{Z};\bm{M};\bm{X}_{b})\geq 0$, where $\bm{X}_{b}$ and
$\bm{X}_{t}$ are the behavior-related $(\bm{s},\bm{a})$-component and task-
related $(\bm{s}^{\prime},r)$-component of the context $\bm{X}$, with
$\bm{X}=(\bm{X}_{t},\bm{X}_{b})$.

For the first graph, $\bm{M}\rightarrow\bm{X}\rightarrow\bm{Z}$ forms a Markov
chain, which satisfies $I(\bm{Z};\bm{M}|\bm{X})=0$. To interpret the second
graph, given an MDP $M\sim\bm{M}$, the state-action component of $\bm{X}$ is
fully captured by the behavior policy $\pi_{\beta}$:
$\bm{s}\sim\mu_{\pi_{\beta}}(\bm{s}),\bm{a}\sim\pi_{\beta}$. We therefore
define it as the behavior-related component, which should be minimally
causally related (dashed lines) to $\bm{M}$ and $\bm{Z}$. Moreover, when
$\bm{X}_{b}$ is given, $\bm{X}_{t}$ is fully characterized by the transition
function $T:(\bm{s},\bm{a})\rightarrow\bm{s}^{\prime}$ and reward function
$R:(\bm{s},\bm{a})\rightarrow r$ of $M$, which should be maximally causally
related (solid lines) to $\bm{M}$ and $\bm{Z}$ and therefore be defined as the
task-related component. Accordingly, we name $I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$
and $I(\bm{Z};\bm{X}_{b})$ the primary and lesser causality in our problem
respectively. With this prior knowledge and assumption, we present the central
theorem of this paper:

######  Theorem 3.3.

Let $\equiv$ denote equality up to a constant,

$\displaystyle\underbrace{I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})}_{\text{primary
causality}}\quad\leq\quad
I(\bm{Z};\bm{M})\quad\leq\underbrace{I(\bm{Z};\bm{X})}_{\text{primary + lesser
causality}}$

holds up to a constant, where

  1. 1.

$\mathcal{L}_{\text{FOCAL}}\equiv-I(\bm{Z};\bm{X})$.

  2. 2.

$\mathcal{L}_{\text{CORRO}}\equiv-I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$.

  3. 3.

$\mathcal{L}_{\text{CSRO}}\geq(\lambda-1)I(\bm{Z};\bm{X})-\lambda
I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$.

###### Proof.

See Appendix B, Theorem B.1. ∎

Theorem 3.3 induces several key observations. Firstly, the FOCAL and CORRO
objective operates as an upper bound and a lower bound of the relevant
information $I(\bm{Z};\bm{M})$ respectively. Since one would like to maximize
$I(\bm{Z};\bm{M})$ according to Definition 3.1, CORRO, which maximizes a lower
bound $I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$, can effectively optimize
$I(\bm{Z};\bm{M})$ with theoretical guarantees. However, FOCAL which maximizes
an upper bound $I(\bm{Z};\bm{X}_{t},\bm{X}_{b})$ is not guaranteed to optimize
$I(\bm{Z};\bm{M})$. Mathematically since
$I(\bm{Z};\bm{X})=I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})+I(\bm{Z};\bm{X}_{b})$,
maximizing the FOCAL objective may also increase the lesser causality
$I(\bm{Z};\bm{X}_{b})$, which is not ideal since it contains spurious
correlation between the task representation $\bm{Z}$ and behavior policy
$\pi_{\beta}$. This explains why FOCAL is less robust to context shift
compared to CORRO.

Secondly, CSRO as a latest COMRL algorithm, inherently optimizes a linear
combination of the FOCAL and CORRO objectives. In the $0\leq\lambda\leq 1$
regime, the CSRO objective becomes a convex interpolation of the upper bound
$I(\bm{Z};\bm{X})$ and the lower bound $I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$ of
$I(\bm{Z};\bm{M})$, which may form a tighter bound on the real objective
$I(\bm{Z};\bm{M})$. This explains the improved performance of CSRO compared to
FOCAL and CORRO.

###  3.3 UNICORN: Implementation of The Information Bottleneck Principle on
COMRL

Given Definition 3.1 of the task representation learning, we now further
generalize the pre-existing COMRL algorithms based on the information
bottleneck principle (Tishby et al., 2000; Tishby and Zaslavsky, 2015).
Namely, finding an optimal task representation $\bm{Z}$ can be formulated as
the minimization of the following Lagrangian

$\displaystyle\mathcal{L}_{IB}=I(\bm{Z};\bm{X})-\beta I(\bm{Z};\bm{M})$ (7)

subject to the Markov chain constraint. The positive Lagrange multiplier
$\beta$ acts as a tradeoff parameter between the complexity/rate of the
representation, $R=I(\bm{Z};\bm{X})$, and the amount of preserved relevant
information, $I_{M}=I(\bm{Z};\bm{M})$. For implementation, the first term is
just FOCAL loss by Theorem 3.3, whereas we propose to approximate the second
term as a convex combination of the FOCAL and CORRO loss, analogous to CSRO:

$\displaystyle I(\bm{Z};\bm{M})\approx\alpha
I(\bm{Z};\bm{X})+(1-\alpha)I(\bm{Z};\bm{X}_{t}|\bm{X}_{b}),$ (8)

Figure 2: Meta-training procedure of UNICORN. Built upon the classic COMRL
methods FOCAL, UNICORN implements the information bottleneck principle by
jointly optimizing the FOCAL loss and a reconstruction loss, using a
variational auto-encoder to obtain robust task representation $\bm{z}$.

where $0\leq\alpha\leq 1$. Substituting Eqn 8 into 7, we may rescale
$\mathcal{L}_{IB}$ as

$\displaystyle\mathcal{L}_{IB}=-\left(I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})+\frac{\alpha\beta-1}{(1-\alpha)\beta}I(\bm{Z};\bm{X})\right).$
(9)

For implementation of $I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$, however, it would
require additional approximations such as generating negative samples for
contrastive learning using a trained model in CORRO, or using a CLUB upper
bound of the mutual information in CSRO. Instead, we propose to rewrite it
according to the chain rule of mutual information (Yeung, 2008):

$\displaystyle I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$
$\displaystyle=I(\bm{X}_{t};\bm{Z},\bm{X}_{b})-I(\bm{X}_{t};\bm{X}_{b})$
$\displaystyle\equiv I(\bm{X}_{t};\bm{Z},\bm{X}_{b}),$ (10)

since $I(\bm{X}_{t};\bm{X}_{b})$ is a constant when $\bm{X}_{t}$ and
$\bm{X}_{b}$ are drawn from a fixed distribution as in offline RL. Moreover,
by definition of mutual information:

$\displaystyle I(\bm{X_{t}};\bm{Z},\bm{X_{b}})$ $\displaystyle=\int
p(\bm{x_{t}},\bm{x_{b}})p(\bm{z}|\bm{x_{t}},\bm{x_{b}})\log\frac{p(\bm{x_{t}}|\bm{z},\bm{x_{b}})}{p(\bm{x_{t}})}$
$\displaystyle\equiv\int
p(\bm{x_{t}},\bm{x_{b}})p(\bm{z}|\bm{x_{t}},\bm{x_{b}})\log
p(\bm{x_{t}}|\bm{z},\bm{x_{b}})$
$\displaystyle\simeq\int_{\bm{x_{t}},\bm{x_{b}},\bm{z}}\underbrace{q_{\bm{\phi}}(\bm{z}|\bm{x_{t}},\bm{x_{b}})}_{\text{encoder}}\underbrace{\log
p_{\bm{\theta}}(\bm{x_{t}}|\bm{z},\bm{x_{b}})}_{\text{decoder}}$
$\displaystyle=\mathbb{E}_{q_{\bm{\phi}}(\bm{z}|\bm{x_{t}},\bm{x_{b}})}\left[\log
p_{\bm{\theta}}(\bm{x_{t}}|\bm{z},\bm{x_{b}})\right],$ (11)

which induces an unbiased estimator of $I(\bm{X}_{t};\bm{Z},\bm{X}_{b})$ by
reconstructing $\bm{X}_{t}$ with a decoder network
$p_{\bm{\theta}}(\cdot|\bm{z},\bm{x_{b}})$ conditioning on $\bm{Z}$ and
$\bm{X_{b}}$, reminiscent of VAE (Kingma and Welling, 2014). We therefore
implement the objective of our proposed algorithm UNICORN as a linear
combination of this reconstruction loss
$\mathcal{L}_{\text{recon}}:=-I(\bm{X_{t}};\bm{Z},\bm{X_{b}})$ and the FOCAL
loss $\mathcal{L}_{\text{FOCAL}}=-I(\bm{Z};\bm{X})$:

$\displaystyle\mathcal{L}_{\text{UNICORN}}$
$\displaystyle:=\mathcal{L}_{IB}=\mathcal{L}_{\text{recon}}+\frac{\alpha\beta-1}{(1-\alpha)\beta}\mathcal{L}_{\text{FOCAL}}$
$\displaystyle=-\left(I(\bm{X}_{t};\bm{Z},\bm{X}_{b})+\frac{\alpha\beta-1}{(1-\alpha)\beta}I(\bm{Z};\bm{X})\right).$
(12)

UNICORN can be interpreted as a generalization of the FOCAL and CORRO
objecitves by effectively striking a balance between the primary causality
$I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$ and the lesser causality
$I(\bm{Z};\bm{X}_{b})$. Figure 2 gives an overview of our framework, which
enjoys three core advantages:

  * •

Robustness and bias-reduction. By optimizing a reconstruction loss
$I(\bm{X}_{t};\bm{Z},\bm{X}_{b})$ equivalent to the CORRO objective
$I(\bm{Z};\bm{X_{t}}|\bm{X_{b}})$, UNICORN mitigates the spurious correlation
in FOCAL, leading to more robust task representation $\bm{Z}$. It is also
verified in Section 4.5 that UNICORN is robust to a wide range of data
qualities. Moreover, the unbiased estimator in Eqn 11 circumvents additional
approximations as in CORRO and CSRO, making UNICORN a bias-reduced and
improved version of previous SOTAs.

  * •

Broader applicability domain. Unlike FOCAL/CORRO which assumes uniform task
distribution, the unbiased estimator in Eqn 11 can be efficiently implemented
by Monte Carlo sampling on offline datasets drawn from any distribution.
Additionally, we will show in Section 5.1 that UNICORN is transferable to
other model architectures such as Decision Transformer (DT).

  * •

General and unified paradigm. Derived from the general information bottleneck
principle, our framework encompasses a series of milestones in COMRL including
FOCAL, CORRO, CSRO and a concurrent method GENTLE (Zhou et al., 2023), which
solely optimizes the reconstruction loss with a similar implementation. A
unified picture of these pre-existing methods being special cases of UNICORN
in the $\alpha$-$\beta$ plane is depicted in Figure 3. A hyper-parameter study
showing the effect of the UNICORN feasible region is presented in Appendix
C.1. Moreover, the decoder $p_{\theta}(\bm{x}_{t}|\bm{z},\bm{x}_{b})$ can
operate as a model of the MDP to generate additional data using customized
latent $\bm{z}$, which is explored in Section 5.2. In summary, UNICORN not
only provides a unified framework of model-free COMRL, it may also be
exploited and lead to novel model-based algorithms.

##  4 Experiments

Our main experiments are organized to address three primary inquiries
regarding the core advantages of UNICORN: (1) How does UNICORN perform on in-
distribution tasks (bias-reduction)? (2) How well can UNICORN generalize to
data collected by OOD behavior policies (robustness)? (3) Can UNICORN
outperform consistently across datasets of different qualities (robustness)?
The performance are measured by the average return across all test tasks.

Figure 3: The hyper-parameter plane of UNICORN. Empirically, we find that
UNICORN works properly only when $(\alpha\beta-1)/(1-\alpha)\beta>0$, which
further constrains the feasible region of UNICORN (shaded) to be bounded by
$\alpha\beta\geq 1$, $\alpha\leq 1$ and $\beta\geq 1$. Intriguingly, pre-
existing methods FOCAL (red) and CORRO/CSRO/GENTLE (blue) constitute the two
boundaries of the feasible region, which can therefore be interpreted as
special cases of UNICORN.

###  4.1 Environments and Datasets

We adopt MuJoCo (Todorov et al., 2012) and MetaWorld (Yu et al., 2020a)
benchmarks to evaluate our method.

HalfCheetah-Dir, HalfCheetah-Vel, Ant-Dir are multi-task MuJoCo benchmarks
where tasks differ in reward function. For HalfCheetah-Dir and Ant-Dir, the
agent is rewarded with high velocity along the goal direction. For
HalfCheetah-Vel, the agent receives an $L_{2}$ penalty with respect to the
target velocity.

Hopper-Param, Walker-Param are multi-task MuJoCo benchmarks where tasks differ
in transition dynamics. The physical parameters of body mass, inertia,
damping, and friction vary for each task.

Reach is a multi-task MetaWorld benchmark where tasks differ in reward
function. The agent controls the end-effector to reach a target position in 3D
space.

Following the protocol of CORRO, we sample 20 training tasks and 20 testing
tasks from the task distribution. We train SAC agent (Haarnoja et al., 2018)
independently for each task and collect offline datasets of mixed quality
using all logged behavior policies.

###  4.2 Baselines

We compare UNICORN with six competitive OMRL algorithms. These baselines are
categorized as context-based, gradient-based or transformer-based methods.

FOCAL, CORRO, CSRO are context-based methods which can be seen as special
cases or approximations of UNICORN, see details in Section 3.

Supervised is a context-based method, directly using actor-critic loss to
train the policy/value networks and task encoder end-to-end. We find it a
competitive baseline across all benchmarks.

MACAW (Mitchell et al., 2021) is a gradient-based method, using supervised
advantage-weighted regression for both the inner and outer loop of meta-
training.

Prompt-DT (Xu et al., 2022) is a transformer-based method, taking context as
the prompt of a Decision Transformer (DT) (Chen et al., 2021; Janner et al.,
2021) to solve OMRL as a conditional sequence modeling problem.

We perform our evaluation in a few-shot manner. For gradient-based methods, we
first update the meta-policy with a batch of context, then evaluate the agent.
For transformer-based methods, we utilize $\bm{z}$ inferred from a batch of
context as the prompt to condition the auto-regressive rollout.

###  4.3 Few-Shot Generalization to In-Distribution Data

Figure 4: Testing returns of UNICORN against baselines on six benchmarks.
Solid curves refer to the mean performance of trials over 6 random seeds, and
the shaded areas characterize the standard deviation of these trials. Table 1:
Average testing returns of UNICORN against baselines on datasets collected by
IID and OOD behavior policies. Each result is averaged by 6 random seeds. The
best is bolded and the second best is underlined.

Algorithm | HalfCheetah-Dir | HalfCheetah-Vel | Ant-Dir | Hopper-Param | Walker-Param | Reach
---|---|---|---|---|---|---
IID | OOD | IID | OOD | IID | OOD | IID | OOD | IID | OOD | IID | OOD
UNICORN | 1312.59±31.72 | 1301.17±25.98 | -21.88±0.65 | -93.57±4.85 | 266.85±14.54 | 235.55±17.69 | 315.90±5.67 | 304.45±10.88 | 407.83±75.58 | 392.48±74.45 | 2774.67±241.04 | 2603.68±182.85
CSRO | 1180.48±228.10 | 458.08±253.00 | -27.98±0.90 | -102.44±5.17 | 276.18±19.37 | 233.33±12.37 | 309.63±6.50 | 301.12±9.99 | 310.47±58.29 | 278.94±65.37 | 2719.50±234.55 | 2800.93±182.06
CORRO | 704.20±449.46 | 244.95±145.83 | -37.11±2.60 | -112.29±1.62 | 147.93±12.85 | 120.52±12.48 | 283.04±7.89 | 271.48±12.79 | 277.07±37.91 | 213.43±48.33 | 2467.83±174.57 | 2321.87±327.19
FOCAL | 1186.03±271.58 | 860.58±252.96 | -22.22±0.88 | -96.72±1.69 | 217.10±29.14 | 173.16±24.37 | 301.84±4.48 | 296.87±12.76 | 308.05±98.42 | 286.34±91.08 | 2423.67±255.95 | 2315.56±303.03
Supervised | 961.62±355.84 | 781.51±429.12 | -23.81±0.71 | -104.31±0.65 | 237.52±38.81 | 202.19±37.56 | 306.54±9.99 | 293.65±7.99 | 256.52±60.39 | 209.64±27.72 | 2489.37±248.05 | 2283.26±205.36
MACAW | 1154.83±10.11 | 449.56±6.54 | -56.04±2.42 | -188.44±0.91 | 26.29±2.80 | 0.18±0.31 | 218.52±6.01 | 205.39±2.47 | 141.25±8.83 | 129.63±5.09 | 2430.67±157.42 | 1728.01±78.98
Prompt-DT | 1176.54±39.64 | -25.10±8.74 | -118.14±65.59 | -248.64±21.38 | 1.00±0.38 | 0.04±0.11 | 234.19±5.35 | 202.51±4.96 | 184.73±8.57 | 155.78±16.64 | 2164.86±85.20 | 1895.51±110.64

Figure 4 illustrates the learning curves of UNICORN vs. all baselines. IID
entries of Table 1 depicts the performance of in-distribution generalization
where the context for each task is collected by the behavior policies trained
on the same task. The results demonstrate that UNICORN almost unanimously
achieves SOTA performance across all benchmarks. The observation that FOCAL,
CORRO, and CSRO achieve asymptotic performance close to UNICORN is also
expected, since they are special cases of UNICORN according to our theory.
Notably, UNICORN exhibits remarkable stability, especially on HalfCheetah-Dir
and Walker-Param, where the other context-based methods suffer a significant
performance decline during the later training process. Although the gradient-
based MACAW and the transformer-based Prompt-DT show faster convergence, their
asymptotic performance leaves much to be desired. Moreover, the average
training time for MACAW is almost three times longer than the context-based
methods.

###  4.4 Few-Shot Generalization to Out-of-Distribution Behavior Policies

Denote by $\\{\pi_{\beta}^{i,t}\\}_{i,t}$ the checkpoints of behavior policies
for all tasks logged during data collection (i.e., training SAC agents), where
$i$ refers to the task index and $t$ represents training iteration. For each
test task $M^{j}$, we use every individual behavior policy in
$\\{\pi_{\beta}^{i,t}\\}_{i,t}$ to collect the OOD contexts and the
conditioned rollouts, as visualized in Figure 1. The OOD testing performance
is measured by averaging returns across all testing tasks, as shown in the OOD
entries of Table 1. While all methods suffer notable decline when facing OOD
context, UNICORN maintains the strongest performance by a large margin.
Another observation is that the context-based methods are in general
significantly more robust compared to the gradient-based MACAW and
transformer-based Prompt-DT, which suggests that context-based paradigm should
be prioritized for research of generalizable OMRL. This also delineates the
storyline of COMRL development from 2021 to date

$\displaystyle\underset{(2021)}{\text{FOCAL}}\rightarrow\underset{(2022)}{\text{CORRO}}\rightarrow\underset{(2023)}{\text{CSRO}}\rightarrow\underset{(2024)}{\text{UNICORN}}$

as a roadmap for pursuing more robust and generalizable task representation
learning for COMRL.

###  4.5 Influence of Data Quality

Table 2: UNICORN vs. baselines on Ant-Dir datasets of various qualities. Each
result is averaged by 6 random seeds. The best is bolded and the second best
is underlined.

Algorithm | Random | Medium | Expert
---|---|---|---
IID | OOD | IID | OOD | IID | OOD
UNICORN | 81.31±17.80 | 62.20±6.15 | 219.86±23.04 | 242.74±10.40 | 278.82±10.46 | 261.83±13.49
CSRO | 1.55±2.65 | -0.20±0.52 | 165.65±9.89 | 198.15±17.35 | 251.77±39.16 | 201.57±45.29
CORRO | 0.87±0.53 | -0.05±0.06 | 8.16±5.02 | -7.01±2.37 | -3.90±9.59 | -13.59±9.06
FOCAL | 67.00±26.33 | 43.65±10.08 | 171.10±84.26 | 187.06±86.45 | 229.02±42.07 | 246.27±20.21
Supervised | 64.97±5.57 | 46.90±11.86 | 149.21±50.21 | 109.85±80.09 | 249.11±33.47 | 214.55±59.56
MACAW | 3.412±0.74 | -0.27±0.25 | 27.69±2.17 | 1.34±0.75 | 87.86±42.69 | 0.64±1.49
Prompt-DT | 0.54±0.29 | 0.02±0.09 | 2.14±4.36 | 0.33±0.57 | 77.94±14.67 | 1.24±2.18

To test whether the UNICORN framework can be applied to different data
distributions, we collect three types of datasets whose size is equal to the
mixed dataset used in Section 4.3 and 4.4: random, medium and expert, which
are characterized by the quality of the behavior policies. Each dataset is
collected by one SAC agent. For random datasets, the SAC agent is randomly
initialized and untrained.

As shown in Table 2, UNICORN demonstrates SOTA performance on all types of
datasets. Notably, despite CSRO having better IID results than UNICORN on Ant-
Dir under the mixed data distribution, UNICORN surpasses CSRO by a large
margin under the narrow data distribution (medium and expert) and random data
distribution.

We observe that CORRO fails on all three datasets. This may be due to the
reliance of CORRO on the negative sample generator. Under narrow data
distribution or poor random data distribution, it might be more challenging to
generate negative samples of good quality. As for Prompt-DT, its performance
improves significantly with increased data quality, which is expected since
the decision transformer adopts a behavior-cloning-like supervised learning
style.

##  5 Discussion

This section presents more empirical evidence on the broader applicability of
the UNICORN framework.

###  5.1 Is UNICORN Model-Agnostic?

Since UNICORN tackles task representation learning from a general information
theoretic perspective, it is in principle model-agnostic. Hence a
straightforward idea is to transfer UNICORN to other model architectures like
DT, which we envision to be a promising backbone for large-scale training of
RL foundation models.

We employ a simple implementation by embedding the task representation vector
$\bm{z}$ as the first token in sequence to prompt the learning and generation
of DT. Instead of the unlearnable prompt used in the original Prompt-DT, we
enforce task representation learning of the prompt by optimizing
$\mathcal{L}_{\text{UNICORN}}$ and $\mathcal{L}_{\text{FOCAL}}$. We name these
DT implementations as UNICORN-DT and FOCAL-DT respectively.

As shown in Table 3 and Figure 8, applying UNICORN and FOCAL on DT results in
significant improvement in both IID and OOD generalization, with UNICORN being
the superior option. Despite the gap between our implementation of UNICORN-DT
and its MLP counterpart in terms of asymptotic performance, we expect UNICORN-
DT to extrapolate favorably in the regime of greater dataset size and model
parameters due to the scaling law of transformer (Kaplan et al., 2020) and DT
(Lee et al., 2022). We leave this verification to future work.

Table 3: DT implementation of COMRL on HalfCheetah-Dir and Hopper-Param. Each
result is averaged by 6 random seeds.

Algorithm | HalfCheetah-Dir | Hopper-Param
---|---|---
IID | OOD | IID | OOD
UNICORN | 1312.59±31.72 | 1301.17±25.98 | 315.90±5.67 | 304.45±10.88
UNICORN-DT | 1233.19±10.42 | 1185.76±42.71 | 303.66±3.51 | 290.88±3.68
FOCAL-DT | 1208.58±33.38 | 652.14±35.86 | 293.25±4.17 | 283.86±4.81
Prompt-DT | 1176.54±39.64 | -25.10±8.74 | 234.19±5.35 | 202.51±4.96

###  5.2 Can UNICORN be Exploited for Model-Based Paradigms?

As alluded in Section 3.3, UNICORN can potentially enable a model-based
paradigm by generating data using the decoder
$p_{\theta}(\bm{x}_{t}|\bm{z},\bm{x}_{b})$ and customized latent $\bm{z}$,
which can be especially useful for generalization to OOD tasks. We implement
this idea on Ant-Dir by sorting the tasks according to the polar angle of
their goal direction ($\theta\in[0,2\pi)$), and construct task-level OOD by
taking 20 consecutive tasks in the middle for training and the rest for
testing (Figure 1). Gaussian noise $\bm{\epsilon}$ is applied to the real task
representation $\bm{z}$ and $(\bm{s},\bm{a},\bm{z}+\bm{\epsilon})$ is then fed
to the decoder to generate $(\bm{s}^{\prime},r)$. The imaginary rollouts
$\\{(\bm{s}_{t},\bm{a}_{t},\bm{s}^{\prime}_{t},r_{t})\\}_{t=1}^{n}$ are used
during the training of RL agent. We adopt the conventional ensemble technique
in model-based RL (Yu et al., 2020b; Zhang et al., 2023) to stabilize the
training process.

As shown in Figure 5, UNICORN-enabled model-based RL is the only method to
achieve positive performance in this extremely challenging scenario where
context shift is induced by OOD tasks.

Figure 5: Testing returns for OOD tasks. The learning curve is averaged by 6
random seeds.

##  6 Conclusion

In this paper, we unify three major context-based offline meta-RL algorithms,
FOCAL, CORRO and CSRO (potentially a lot more) into a single framework based
on the information bottleneck principle. Our theory offers valuable insight on
how these methods are inherently connected and incrementally evolved in
pursuit of better generalization against context shift. Based on the proposed
framework, we arrive at a novel algorithm called UNICORN which is demonstrated
to be remarkably robust and broadly applicable through extensive experiments.
We believe our study offers potential for new implementations, optimality
bounds and algorithms for both fully offline and offline-to-online COMRL
paradigms. We leave these directions for future work.

## References

  * Beck et al. [2023] J. Beck, R. Vuorio, E. Z. Liu, Z. Xiong, L. Zintgraf, C. Finn, and S. Whiteson.  A survey of meta-reinforcement learning.  _arXiv preprint arXiv:2301.08028_ , 2023.
  * Bommasani et al. [2021] R. Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein, J. Bohg, A. Bosselut, E. Brunskill, et al.  On the opportunities and risks of foundation models.  _arXiv preprint arXiv:2108.07258_ , 2021.
  * Chen et al. [2021] L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch.  Decision transformer: Reinforcement learning via sequence modeling.  _Advances in neural information processing systems_ , 34:15084–15097, 2021.
  * Cheng et al. [2020] P. Cheng, W. Hao, S. Dai, J. Liu, Z. Gan, and L. Carin.  Club: A contrastive log-ratio upper bound of mutual information.  In _International conference on machine learning_ , pages 1779–1788. PMLR, 2020.
  * Dorfman et al. [2021] R. Dorfman, I. Shenfeld, and A. Tamar.  Offline meta reinforcement learning–identifiability challenges and effective data collection strategies.  _Advances in Neural Information Processing Systems_ , 34:4607–4618, 2021.
  * Fakoor et al. [2020] R. Fakoor, P. Chaudhari, S. Soatto, and A. J. Smola.  Meta-q-learning.  In _International Conference on Learning Representations_ , 2020.
  * Feng et al. [2022] L. Feng, P. Nouri, A. Muni, Y. Bengio, and P.-L. Bacon.  Designing biological sequences via meta-reinforcement learning and bayesian optimization.  _arXiv preprint arXiv:2209.06259_ , 2022.
  * Finn et al. [2017] C. Finn, P. Abbeel, and S. Levine.  Model-agnostic meta-learning for fast adaptation of deep networks.  In _International conference on machine learning_ , pages 1126–1135. PMLR, 2017.
  * Gao et al. [2022] C. Gao, K. Xu, K. Zhou, L. Li, X. Wang, B. Yuan, and P. Zhao.  Value penalized q-learning for recommender systems.  In _Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval_ , pages 2008–2012, 2022\.
  * Gao et al. [2023] Y. Gao, R. Zhang, J. Guo, F. Wu, Q. Yi, S. Peng, S. Lan, R. Chen, Z. Du, X. Hu, et al.  Context shift reduction for offline meta-reinforcement learning.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.
  * Geirhos et al. [2020] R. Geirhos, J.-H. Jacobsen, C. Michaelis, R. Zemel, W. Brendel, M. Bethge, and F. A. Wichmann.  Shortcut learning in deep neural networks.  _Nature Machine Intelligence_ , 2(11):665–673, 2020.
  * Gottesman et al. [2019] O. Gottesman, F. Johansson, M. Komorowski, A. Faisal, D. Sontag, F. Doshi-Velez, and L. A. Celi.  Guidelines for reinforcement learning in healthcare.  _Nature medicine_ , 25(1):16–18, 2019.
  * Gulrajani et al. [2017] I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. C. Courville.  Improved training of wasserstein gans.  _Advances in neural information processing systems_ , 30, 2017.
  * Haarnoja et al. [2018] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _International conference on machine learning_ , pages 1861–1870. PMLR, 2018.
  * Hallak et al. [2015] A. Hallak, D. Di Castro, and S. Mannor.  Contextual markov decision processes.  _arXiv preprint arXiv:1502.02259_ , 2015.
  * Hu et al. [2021] E. J. Hu, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, W. Chen, et al.  Lora: Low-rank adaptation of large language models.  In _International Conference on Learning Representations_ , 2021.
  * Janner et al. [2021] M. Janner, Q. Li, and S. Levine.  Offline reinforcement learning as one big sequence modeling problem.  _Advances in neural information processing systems_ , 34:1273–1286, 2021.
  * Jaques et al. [2019] N. Jaques, A. Ghandeharioun, J. H. Shen, C. Ferguson, A. Lapedriza, N. Jones, S. Gu, and R. Picard.  Way off-policy batch deep reinforcement learning of implicit human preferences in dialog.  _arXiv preprint arXiv:1907.00456_ , 2019.
  * Kaelbling et al. [1998] L. P. Kaelbling, M. L. Littman, and A. R. Cassandra.  Planning and acting in partially observable stochastic domains.  _Artificial intelligence_ , 101(1-2):99–134, 1998\.
  * Kaplan et al. [2020] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_ , 2020.
  * Kingma and Welling [2014] D. P. Kingma and M. Welling.  Auto-encoding variational bayes.  In Y. Bengio and Y. LeCun, editors, _2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings_ , 2014.  URL http://arxiv.org/abs/1312.6114.
  * Kumar et al. [2019] A. Kumar, J. Fu, M. Soh, G. Tucker, and S. Levine.  Stabilizing off-policy q-learning via bootstrapping error reduction.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Kumar et al. [2021] A. Kumar, Z. Fu, D. Pathak, and J. Malik.  Rma: Rapid motor adaptation for legged robots.  _arXiv preprint arXiv:2107.04034_ , 2021.
  * Lee et al. [2022] K.-H. Lee, O. Nachum, M. S. Yang, L. Lee, D. Freeman, S. Guadarrama, I. Fischer, W. Xu, E. Jang, H. Michalewski, et al.  Multi-game decision transformers.  _Advances in Neural Information Processing Systems_ , 35:27921–27936, 2022.
  * Lester et al. [2021] B. Lester, R. Al-Rfou, and N. Constant.  The power of scale for parameter-efficient prompt tuning.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pages 3045–3059, 2021.
  * Levine et al. [2020] S. Levine, A. Kumar, G. Tucker, and J. Fu.  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  _arXiv preprint arXiv:2005.01643_ , 2020.
  * Li et al. [2020] J. Li, Q. Vuong, S. Liu, M. Liu, K. Ciosek, H. Christensen, and H. Su.  Multi-task batch reinforcement learning with metric learning.  _Advances in Neural Information Processing Systems_ , 33:6197–6210, 2020.
  * Li et al. [2021a] L. Li, Y. Huang, M. Chen, S. Luo, D. Luo, and J. Huang.  Provably improved context-based offline meta-rl with attention and contrastive learning.  _arXiv preprint arXiv:2102.10774_ , 2021a.
  * Li et al. [2021b] L. Li, R. Yang, and D. Luo.  Focal: Efficient fully-offline meta-reinforcement learning via distance metric learning and behavior regularization.  In _International Conference on Learning Representations_ , 2021b.
  * Li and Liang [2021] X. L. Li and P. Liang.  Prefix-tuning: Optimizing continuous prompts for generation.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , pages 4582–4597, 2021.
  * Mishra et al. [2017] N. Mishra, M. Rohaninejad, X. Chen, and P. Abbeel.  Meta-learning with temporal convolutions.  _arXiv preprint arXiv:1707.03141_ , 2(7):23, 2017\.
  * Mitchell et al. [2021] E. Mitchell, R. Rafailov, X. B. Peng, S. Levine, and C. Finn.  Offline meta-reinforcement learning with advantage weighting.  In _International Conference on Machine Learning_ , pages 7780–7791. PMLR, 2021.
  * Oord et al. [2018] A. v. d. Oord, Y. Li, and O. Vinyals.  Representation learning with contrastive predictive coding.  _arXiv preprint arXiv:1807.03748_ , 2018.
  * Pong et al. [2022] V. H. Pong, A. V. Nair, L. M. Smith, C. Huang, and S. Levine.  Offline meta-reinforcement learning with online self-supervision.  In _International Conference on Machine Learning_ , pages 17811–17829. PMLR, 2022.
  * Raffel et al. [2020] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _The Journal of Machine Learning Research_ , 21(1):5485–5551, 2020.
  * Rakelly et al. [2019] K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen.  Efficient off-policy meta-reinforcement learning via probabilistic context variables.  In _International conference on machine learning_ , pages 5331–5340. PMLR, 2019.
  * Reed et al. [2022] S. Reed, K. Zolna, E. Parisotto, S. G. Colmenarejo, A. Novikov, G. Barth-maron, M. Giménez, Y. Sulsky, J. Kay, J. T. Springenberg, et al.  A generalist agent.  _Transactions on Machine Learning Research_ , 2022.
  * Schaul et al. [2015] T. Schaul, D. Horgan, K. Gregor, and D. Silver.  Universal value function approximators.  In _International conference on machine learning_ , pages 1312–1320. PMLR, 2015.
  * Shalev-Shwartz et al. [2016] S. Shalev-Shwartz, S. Shammah, and A. Shashua.  Safe, multi-agent, reinforcement learning for autonomous driving.  _arXiv preprint arXiv:1610.03295_ , 2016.
  * Tishby and Zaslavsky [2015] N. Tishby and N. Zaslavsky.  Deep learning and the information bottleneck principle.  In _2015 ieee information theory workshop (itw)_ , pages 1–5. IEEE, 2015.
  * Tishby et al. [1999] N. Tishby, F. Pereira, W. Bialek, B. Hajek, and R. Sreenivas.  Proceedings of the 37th annual allerton conference on communication, control and computing.  In _Proceedings ao the 37-th Annual Allerton Conference on Communication, Control and Computing, chapter The information bottleneck method_ , pages 368–377, 1999.
  * Tishby et al. [2000] N. Tishby, F. C. Pereira, and W. Bialek.  The information bottleneck method.  _arXiv preprint physics/0004057_ , 2000.
  * Todorov et al. [2012] E. Todorov, T. Erez, and Y. Tassa.  Mujoco: A physics engine for model-based control.  In _2012 IEEE/RSJ international conference on intelligent robots and systems_ , pages 5026–5033. IEEE, 2012.
  * Van der Maaten and Hinton [2008] L. Van der Maaten and G. Hinton.  Visualizing data using t-sne.  _Journal of machine learning research_ , 9(11), 2008.
  * Wang et al. [2016] J. X. Wang, Z. Kurth-Nelson, D. Tirumala, H. Soyer, J. Z. Leibo, R. Munos, C. Blundell, D. Kumaran, and M. Botvinick.  Learning to reinforcement learn.  _arXiv preprint arXiv:1611.05763_ , 2016.
  * Wu et al. [2019] Y. Wu, G. Tucker, and O. Nachum.  Behavior regularized offline reinforcement learning.  _arXiv preprint arXiv:1911.11361_ , 2019.
  * Xu et al. [2022] M. Xu, Y. Shen, S. Zhang, Y. Lu, D. Zhao, J. Tenenbaum, and C. Gan.  Prompting decision transformer for few-shot policy generalization.  In _international conference on machine learning_ , pages 24631–24645. PMLR, 2022.
  * Yeung [2008] R. W. Yeung.  _Information theory and network coding_.  Springer Science & Business Media, 2008.
  * Yu et al. [2020a] T. Yu, D. Quillen, Z. He, R. Julian, K. Hausman, C. Finn, and S. Levine.  Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning.  In _Conference on robot learning_ , pages 1094–1100. PMLR, 2020a.
  * Yu et al. [2020b] T. Yu, G. Thomas, L. Yu, S. Ermon, J. Y. Zou, S. Levine, C. Finn, and T. Ma.  Mopo: Model-based offline policy optimization.  _Advances in Neural Information Processing Systems_ , 33:14129–14142, 2020b.
  * Yuan and Lu [2022] H. Yuan and Z. Lu.  Robust task representations for offline meta-reinforcement learning via contrastive learning.  In _International Conference on Machine Learning_ , pages 25747–25759. PMLR, 2022.
  * Zhang et al. [2023] H. Zhang, H. Yu, J. Zhao, D. Zhang, C. Huang, H. Zhou, X. Zhang, and C. Ye.  How to fine-tune the model: Unified model shift and model bias policy optimization.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.
  * Zhao et al. [2022a] M. Zhao, P. Abbeel, and S. James.  On the effectiveness of fine-tuning versus meta-reinforcement learning.  _Advances in Neural Information Processing Systems_ , 35:26519–26531, 2022a.
  * Zhao et al. [2022b] T. Z. Zhao, J. Luo, O. Sushkov, R. Pevceviciute, N. Heess, J. Scholz, S. Schaal, and S. Levine.  Offline meta-reinforcement learning for industrial insertion.  In _2022 International Conference on Robotics and Automation (ICRA)_ , pages 6386–6393. IEEE, 2022b.
  * Zhou et al. [2023] R. Zhou, C.-X. Gao, Z. Zhang, and Y. Yu.  Generalizable task representation learning for offline meta-reinforcement learning with data limitations.  _arXiv preprint arXiv:2312.15909_ , 2023.
  * Zintgraf et al. [2019] L. Zintgraf, K. Shiarli, V. Kurin, K. Hofmann, and S. Whiteson.  Fast context adaptation via meta-learning.  In _International Conference on Machine Learning_ , pages 7693–7702. PMLR, 2019.

##  Appendix A Pseudo-Code

Algorithm 1 UNICORN Meta-training

0: Offline Datasets $\bm{X}=\\{X^{i}\\}_{i=1}^{N_{train}}$, training tasks
$\bm{M}=\\{M^{i}\\}_{i=1}^{N_{train}}$, initialized learned policy
$\pi_{\omega}$, Q function $Q_{\psi}$, context encoder $q_{\phi}$, decoder
$p_{\theta}$, hyper-parameter $\alpha$, $\beta$, task batch size $N_{tbatch}$

1: while not done do

2: Sample task batch $\\{M^{j}\\}_{j=1}^{N_{tbatch}}\sim\bm{M}$ and the
corresponding replay buffer
$\mathcal{R}=\\{X^{j}\\}_{j=1}^{N_{tbatch}}\sim\bm{X}$

3: for step in each iter do

4: // Train the context encoder and decoder

5: Sample context $\\{\bm{c}^{j}\\}_{j=1}^{N_{tbatch}}\sim\mathcal{R}$

6: Obtain task embeddings $\\{\bm{z}^{j}\sim
q_{\phi}(\bm{z}|\bm{c}^{j})\\}_{j=1}^{N_{tbatch}}$

7: Compute
$\mathcal{L}_{\text{UNICORN}}=\mathcal{L}_{\text{recon}}+\frac{\alpha\beta-1}{(1-\alpha)\beta}\mathcal{L}_{\text{FOCAL}}$

8: Update $\phi,\theta$ to minimize $\mathcal{L}_{\text{UNICORN}}$

9: // Train the actor and critic

10: Detach the task embeddings $\\{\bm{z}^{j}\\}_{j=1}^{N_{tbatch}}$

11: Sample training data $\\{\bm{d}^{j}\\}_{j=1}^{N_{tbatch}}\sim\mathcal{R}$

12: Compute $\mathcal{L}_{\text{actor}},\mathcal{L}_{\text{critic}}$

13: Update $\omega,\psi$ to minimize
$\mathcal{L}_{\text{actor}},\mathcal{L}_{\text{critic}}$

14: end for

15: end while

Algorithm 2 UNICORN Meta-testing

0: Offline Datasets $\bm{X}=\\{X^{i}\\}_{i=1}^{N_{test}}$, testing tasks
$\bm{M}=\\{M^{i}\\}_{i=1}^{N_{test}}$, learned policy $\pi_{\omega}$, context
encoder $q_{\phi}$

1: for each $M^{i}$ do

2: Sample context $\bm{c}^{i}\sim X^{i}$

3: Obtain task embedding $\bm{z}^{i}\sim q_{\phi}(\bm{z}|\bm{c}^{i})$

4: Rollout policy $\pi_{\omega}(\bm{a}|\bm{s},\bm{z}^{i})$ for evaluation

5: end for

##  Appendix B Proofs

######  Theorem B.1.

Let $\equiv$ denote equality up to a constant,

$\displaystyle\underbrace{I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})}_{\text{primary
causality}}\quad\leq\quad
I(\bm{Z};\bm{M})\quad\leq\underbrace{I(\bm{Z};\bm{X})}_{\text{primary + lesser
causality}}$

holds up to a constant, where

  1. 1.

$\mathcal{L}_{\text{FOCAL}}\equiv-I(\bm{Z};\bm{X})$.

  2. 2.

$\mathcal{L}_{\text{CORRO}}\equiv-I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$.

  3. 3.

$\mathcal{L}_{\text{CSRO}}\geq(\lambda-1)I(\bm{Z};\bm{X})-\lambda
I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$.

###### Proof.

Assuming 2,
$I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})\equiv-\mathcal{L}_{\text{CORRO}}\leq
I(\bm{Z};\bm{M})$ is shown in the Theorem 4.1. of CORRO [Yuan and Lu, 2022]
and GENTLE [Zhou et al., 2023] with strong assumptions. We hereby present a
proof with minimal assumptions. Given $I(\bm{Z};\bm{M};\bm{X}_{b})\geq 0$ in
Assumption 3.2, we have

$\displaystyle I(\bm{Z};\bm{M})$
$\displaystyle=I(\bm{Z};\bm{M}|\bm{X}_{b})+I(\bm{Z};\bm{M};\bm{X}_{b})$
$\displaystyle\geq I(\bm{Z};\bm{M}|\bm{X}_{b})$
$\displaystyle=I(\bm{M};\bm{Z},\bm{X}_{t}|\bm{X}_{b})-I(\bm{M};\bm{X}_{t}|\bm{Z},\bm{X}_{b})$
$\displaystyle=\underbrace{I(\bm{Z};\bm{M}|\bm{X}_{t},\bm{X}_{b})}_{0\text{ by
}\bm{M}\rightarrow\bm{X}\rightarrow\bm{Z}}+I(\bm{M};\bm{X}_{t}|\bm{X}_{b})-I(\bm{M};\bm{X}_{t}|\bm{Z},\bm{X}_{b})$
$\displaystyle=I(\bm{M};\bm{X}_{t}|\bm{X}_{b})-I(\bm{M};\bm{X}_{t}|\bm{Z},\bm{X}_{b})$
$\displaystyle=I(\bm{M};\bm{X}_{t}|\bm{X}_{b})-H(\bm{X}_{t}|\bm{Z},\bm{X}_{b})+H(\bm{X}_{t}|\bm{M},\bm{Z},\bm{X}_{b})$
$\displaystyle\geq
I(\bm{M};\bm{X}_{t}|\bm{X}_{b})-H(\bm{X}_{t}|\bm{Z},\bm{X}_{b})$
$\displaystyle=\underbrace{I(\bm{M};\bm{X}_{t}|\bm{X}_{b})}_{\text{const}}-\underbrace{H(\bm{X}_{t})}_{\text{const}}+I(\bm{X}_{t};\bm{Z},\bm{X}_{b})$
$\displaystyle\equiv I(\bm{X}_{t};\bm{Z},\bm{X}_{b})$
$\displaystyle=I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})+\underbrace{I(\bm{X}_{t};\bm{X}_{b})}_{\text{const}}$
$\displaystyle\equiv I(\bm{Z};\bm{X}_{t}|\bm{X}_{b}),$

as desired. The third and fourth lines are obtained by direct application of
the chain rule of mutual information [Yeung, 2008]. All mutual information
terms without $\bm{Z}$ are held constant in the fully offline scenario.

$I(\bm{Z};\bm{M})\leq I(\bm{Z};\bm{X})$ is a direct consequence of the Data
Processing Inequality [Yeung, 2008] and the Markov chain
$\bm{M}\rightarrow\bm{X}\rightarrow\bm{Z}$ in Assumption 3.2. We now prove the
three claims regarding FOCAL, CORRO and CSRO:

  1. 1.

$\mathcal{L}_{\text{FOCAL}}\equiv-I(\bm{Z};\bm{X})$.

$\displaystyle I(\bm{Z};\bm{X})$
$\displaystyle:=\mathbb{E}_{\bm{x},\bm{z}}\left[\log\left(\frac{p(\bm{z},\bm{x})}{p(\bm{z})p(\bm{x})}\right)\right]$
$\displaystyle=\mathbb{E}_{\bm{x},\bm{z}}\left[\log\left(\frac{1}{\frac{p(\bm{z})}{p(\bm{z}|\bm{x})}|\mathcal{M}|}\right)\right]+\log(|\mathcal{M}|)$
$\displaystyle=\mathbb{E}_{\bm{x},\bm{z}}\left[\log\left(\frac{1}{\frac{p(\bm{z})}{p(\bm{z}|\bm{x})}|\mathcal{M}|\mathbb{E}_{\bm{x}}\left[\frac{p(\bm{z}|\bm{x})}{p(\bm{z})}\right]}\right)\right]+\log(|\mathcal{M}|)$
$\displaystyle\approx\mathbb{E}_{\bm{x},\bm{z}}\left[\log\left(\frac{1}{\frac{p(\bm{z})}{p(\bm{z}|\bm{x})}\sum_{M^{i}\in\mathcal{M}}\mathbb{E}_{\bm{x}^{i}\sim
X^{i}}\left[\frac{p(\bm{z}|\bm{x}^{i})}{p(\bm{z})}\right]}\right)\right]+\log(|\mathcal{M}|)$
$\displaystyle=\mathbb{E}_{\bm{x},\bm{z}}\left[\log\left(\frac{\frac{p(\bm{z}|\bm{x})}{p(\bm{z})}}{\sum_{M^{*}\in\mathcal{M}}\mathbb{E}_{\bm{x}^{i}\sim
X^{i}}\left[\frac{p(\bm{z}|\bm{x}^{i})}{p(\bm{z})}\right]}\right)\right]+\log(|\mathcal{M}|)$
$\displaystyle=\mathbb{E}_{\bm{x},\bm{z}}\left[\log\left(\frac{h(\bm{x},\bm{z})}{\sum_{M^{i}\in\mathcal{M}}h(\bm{x}^{i},\bm{z})}\right)\right]+\log(|\mathcal{M}|)$

The first term on RHS is precisely the supervised contrastive learning
objective introduced by a variant of FOCAL [Li et al., 2021a], which is
equivalent to the negative distance metric learning loss
$\mathcal{L}_{\text{FOCAL}}$ with the effect of pushing away embeddings of
different tasks while pulling together those from the same task. Therefore we
have

$\displaystyle I(\bm{Z};\bm{X})$
$\displaystyle=\mathbb{E}_{\bm{x},\bm{z}}\left[\log\left(\frac{h(\bm{x},\bm{z})}{\sum_{M^{i}\in\mathcal{M}}h(\bm{x}^{i},\bm{z})}\right)\right]+\text{const}$
$\displaystyle\equiv-\mathcal{L}_{\text{FOCAL}}.$

  2. 2.

$\mathcal{L}_{\text{CORRO}}\equiv-I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$.

$\displaystyle I(\bm{Z};\bm{X_{t}}|\bm{X_{b}})$
$\displaystyle:=\mathbb{E}_{\bm{x}_{t},\bm{x}_{b},\bm{z}}\left[\log\left(\frac{p(\bm{z},\bm{x_{t}}|\bm{x_{b}})}{p(\bm{z}|\bm{\bm{x_{b}}})p(\bm{x_{t}}|\bm{\bm{x_{b}}})}\right)\right]$
$\displaystyle=\mathbb{E}_{\bm{x_{t}},\bm{\bm{x_{b}}},\bm{z}}\left[\log\left(\frac{1}{\frac{p(\bm{z}|\bm{\bm{x_{b}}})}{p(\bm{z}|\bm{x_{t}},\bm{\bm{x_{b}}})}|\mathcal{M}|}\right)\right]+\log(|\mathcal{M}|)$
$\displaystyle=\mathbb{E}_{\bm{x_{t}},\bm{\bm{x_{b}}},\bm{z}}\left[\log\left(\frac{1}{\frac{p(\bm{z}|\bm{x}_{b})}{p(\bm{z}|\bm{X})}|\mathcal{M}|\mathbb{E}_{M^{*}\in\mathcal{M}}\left[\frac{p(\bm{z}|\bm{x}^{*}_{t},\bm{\bm{x_{b}}})}{p(\bm{z}|\bm{\bm{x_{b}}})}\right]}\right)\right]+\log(|\mathcal{M}|)$
$\displaystyle\approx\mathbb{E}_{\bm{x_{t}},\bm{\bm{x_{b}}},\bm{z}}\left[\log\left(\frac{1}{\frac{p(\bm{z}|\bm{\bm{x_{b}}})}{p(\bm{z}|\bm{x_{t}},\bm{\bm{x_{b}}})}\sum_{M^{*}\in\mathcal{M}}\frac{p(\bm{z}|\bm{x}^{*}_{t},\bm{\bm{x_{b}}})}{p(\bm{z}|\bm{\bm{x_{b}}})}}\right)\right]+\log(|\mathcal{M}|)$
$\displaystyle=\mathbb{E}_{\bm{x_{t}},\bm{\bm{x_{b}}},\bm{z}}\left[\log\left(\frac{\frac{p(\bm{z}|\bm{x_{t}},\bm{\bm{x_{b}}})}{p(\bm{z}|\bm{\bm{x_{b}}})}}{\sum_{M^{*}\in\mathcal{M}}\frac{p(\bm{z}|\bm{x}^{*}_{t},\bm{\bm{x_{b}}})}{p(\bm{z}|\bm{\bm{x_{b}}})}}\right)\right]+\log(|\mathcal{M}|)$
$\displaystyle=\mathbb{E}_{\bm{x_{t}},\bm{\bm{x_{b}}},\bm{z}}\left[\log\left(\frac{\frac{p(\bm{z}|\bm{x_{t}},\bm{\bm{x_{b}}})}{p(z)}}{\sum_{M^{*}\in\mathcal{M}}\frac{p(\bm{z}|\bm{x}^{*}_{t},\bm{\bm{x_{b}}})}{p(z)}}\right)\right]+\log(|\mathcal{M}|)$
$\displaystyle=\mathbb{E}_{\bm{x},\bm{z}}\left[\log\left(\frac{h(\bm{x},\bm{z})}{\sum_{M^{*}\in\mathcal{M}}h(\bm{x}^{*},\bm{z})}\right)\right]+\log(|\mathcal{M}|)$
$\displaystyle\equiv-\mathcal{L}_{\text{CORRO}}.$

  3. 3.

$\mathcal{L}_{\text{CSRO}}\geq(\lambda-1)I(\bm{Z};\bm{X})-\lambda
I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$.

The CLUB loss in Eqn 6 operates as a upper bound of $I(\bm{Z};\bm{X}_{b})$
[Cheng et al., 2020]. By Eqn 5 we have

$\displaystyle\mathcal{L}_{\text{CSRO}}$
$\displaystyle=\mathcal{L}_{\text{FOCAL}}+\lambda L_{\text{CLUB}}$
$\displaystyle\geq\mathcal{L}_{\text{FOCAL}}+\lambda I(\bm{Z};\bm{X}_{b})$
$\displaystyle=\mathcal{L}_{\text{FOCAL}}+\lambda\underbrace{(I(\bm{Z};\bm{X})-I(\bm{Z};\bm{X}_{t}|\bm{X}_{b}))}_{\text{chain
rule of mutual information}}$
$\displaystyle\overset{1}{=}-I(\bm{Z};\bm{X})+\lambda(I(\bm{Z};\bm{X})-I(\bm{Z};\bm{X}_{t}|\bm{X}_{b}))$
$\displaystyle=(\lambda-1)I(\bm{Z};\bm{X})-\lambda
I(\bm{Z};\bm{X}_{t}|\bm{X}_{b}).$

∎

##  Appendix C Further Experiments

Figure 6: Hyper-parameter study results in all environments, each result is
averaged over 6 random seeds.

###  C.1 Hyper-parameter Study of $\frac{\alpha\beta-1}{(1-\alpha)\beta}$

We experiment with $\frac{\alpha\beta-1}{(1-\alpha)\beta}$ from 0 to 10, as
shown in Figure 6. We observe that UNICORN has an optimal
$\frac{\alpha\beta-1}{(1-\alpha)\beta}$ in all environments. This observation
suggests that it is insufficient to use either FOCAL loss or reconstruction
loss alone, emphasizing the necessity of their combination proposed in Eqn 7
and 8.

Intuitively, for tasks differ in transition dynamics, the lesser causality
$I(\bm{Z};\bm{X}_{b})$ should be higher since the state distribution
$\mu_{\pi}(\bm{s})$ vary across tasks even with the same behavior policy
$\pi$. This indicates that the optimal proportion of FOCAL loss in Eqn 12
should be greater on these environments, as it maximizes $I(\bm{Z};\bm{X})$
along with the primary causality $I(\bm{Z};\bm{X}_{t}|\bm{X}_{b})$. Such
thought experiment is verified in this study. We record the optimal
$\frac{\alpha\beta-1}{(1-\alpha)\beta}$ of all benchmarks in Table 5,
demonstrating that for tasks differ in reward function (Ant-Dir, HalfCheetah-
Dir) and transition dynamics (Hopper-Param, Walker-Param), the optimal value
of $\frac{\alpha\beta-1}{(1-\alpha)\beta}$ is around 0.1 and 1 respectively.
The fact that the optimal $\frac{\alpha\beta-1}{(1-\alpha)\beta}$ does not
vanish for Ant-Dir and HalfCheetah-Dir suggests that the lesser causality
$I(\bm{Z};\bm{X}_{b})$ should be nonzero even when all tasks share the same
transition dynamics. We conjecture this is partially due to the fact that the
behavior policy we use for each task is also trained on that task, which leads
to task-dependent state distribution $\mu_{\pi}(\bm{s})$. However, this
applies to the IID setting only.

###  C.2 UNICORN-0: A Label-free Version

When $\frac{\alpha\beta-1}{(1-\alpha)\beta}$ equals zero, UNICORN task
representation learning reduces to the minimization of the reconstruction loss
only, which becomes a label-free algorithm (i.e., it does not require the
knowledge of task identities/labels to optimize, as opposed to the contrastive
learning in previous methods). We name this special case UNICORN-0, which is
equivalent to a concurrent method GENTLE [Zhou et al., 2023]. For a fair
comparison, we use a variant of a Bayesian OMRL method BOReL [Dorfman et al.,
2021] that does not utilize oracle reward functions, as a label-free baseline.
At test time, we modify BOReL to use offline datasets rather than the online
data collected by interacting with the environment to infer the task
information111Experiment shows that the usage of offline dataset rather than
online data has little effect on the results of BOReL.. As shown in Table 4,
UNICORN-0 is also competitive with BOReL.

Table 4: UNICORN-0 compared to another label-free COMRL method, BOReL, on Ant-
Dir.

Algorithm | Ant-Dir
---|---
IID | OOD
UNICORN-0 | 219.88±16.24 | 199.86±9.26
BOReL | 206.4±17.93 | 186.55±9.54

###  C.3 Visualization of Task Representations

For better interpretation, we visualize the task representations by 2D
projection of the embedding vectors via t-SNE [Van der Maaten and Hinton,
2008]. For each test task, we generate a trajectory using each policy in
$\\{\pi_{\beta}^{i,t}\\}_{i,t}$ and infer their task representations $\bm{z}$.
This setting aligns with our OOD experiments in the main text.

Since the UNICORN objective operates as a convex combination of the
reconstruction loss (UNICORN-0) and the FOCAL loss, we compare representations
of UNICORN to these two extremes. As shown in Figure 7, compared to UNICORN-0,
UNICORN provides evidently more distinguishable task representations for OOD
data, emphasizing the necessity of joint optimization with the FOCAL loss. On
the other hand, despite the seemingly higher quality of task representations
of FOCAL, the performance of FOCAL is much worse than UNICORN. We speculate
that since FOCAL blindly separates embeddings from different tasks, it may
fail to capture shared structure between similar tasks, whereas UNICORN is
able to do so by generative modeling via the reconstruction loss.

Figure 7: The 2D projection of the learned task representation space in Ant-
Dir. Points are uniformly sampled from out-of-distribution data. Tasks of
given goals from 0 to 6 are mapped to rainbow colors, ranging from purple to
red.

###  C.4 More on UNICORN-DT

We provide the learning curves of the results in Section 5.1 here. As shown in
Figure 8, UNICORN-DT demonstrates both the higher convergence speed and the
higher asymptotic performance compared to the vanilla Prompt-DT.

Figure 8: Test returns of UNICORN-DT against UNICORN and Prompt-DT on 2
benchmarks. The learning curve is averaged by 6 random seeds.

##  Appendix D Experimental Details

Table 5 lists the necessary hyper-parameters that we used to produce the
experimental results.

Table 5: Configurations and hyper-parameters used in the training process.

Configurations | Ant-Dir | HalfCheetah-Dir | HalfCheetah-Vel | Hopper-Param | Walker-Param | Reach
---|---|---|---|---|---|---
dataset size | 1e5 | 2e5 | 2e5 | 3e5 | 4.5e5 | 1e5
task representation dimension | 5 | 5 | 5 | 40 | 32 | 2
optimal $\frac{\alpha\beta-1}{(1-\alpha)\beta}$ | 0.15 | 0.01 | 0.15 | 1.5 | 0.7 | 0.3
training steps | 200k
task batch size | 16
RL batch size | 256
context training size | 200 | 500
learning rate | 3e-4
RL network width | 256
RL network depth | 3
encoder width | 64 | 128 | 64