Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning
Authors: Lanqing Li
Hai Zhang
Xinyu Zhang
Shatong Zhu
Yang YU
Junqiao Zhao
Pheng-Ann Heng
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Offline Meta-RL
Information Theory
Primary area: reinforcement_learning
TL;DR: We propose a novel information theoretic framework of the context-based offline meta-RL paradigm, which unifies several mainstream methods and leads to two robust algorithm implementations.

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 framework. Most importantly, we show that the pre-existing COMRL algorithms are essentially optimizing the same mutual information objective between the task variable $M$ and its latent representation $Z$ by implementing various approximate bounds. Such theoretical insight offers ample design freedom for novel algorithms. As demonstrations, we propose a supervised and a self-supervised implementation of $I(Z; M)$, and empirically show that the corresponding optimization algorithms exhibit remarkable generalization across a broad spectrum of RL benchmarks, context shift scenarios, data qualities and deep learning architectures. This work lays the information theoretic foundation for COMRL methods, leading to a better understanding of task representation learning in the context of reinforcement learning. Given its
generality, we envision our framework as a promising offline pre-training paradigm of foundation models for decision making.

Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning
Langing Li Hai Zhang': Xinyu Zhang", Shatong Zhu', Yang Yu',
Junqiao Zhao Pheng-Ann Heng
Zhejiang Lab, The Chinese University of Hong Kong,
Tongji University, Stony Brook University
langingli 1at.gngg1119990mm....mmmmmmmm com, (thanghai12138, zhushatong, zhaojunqiao (oaggiaa)@toogjigi edu.. en, zhang1469cs. stw.ssbooo..ooooooooo.oooooooo.. edu, (yangyu, pheng) cuhk. edu.hk
Abstract
As marriage between offline RL and meta-RL,, the advent of offline meta- reinforcement learning (OMRL) has shown great promise in enabling RI. agents multi-task and quickly adapt while acquiring knowledge safely. Among which, context-based OMRL (COMRL) as popular paradigm, aims to learn 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 seem- ingly independent methodologies into unified framework. Most importantly, we show that the pre-existing COMRL algorithms are essentially optimizing the same mutual information objective between the task variable M and its latent representation Z by implementing various approximate bounds. Such theoretical insight offers ample design freedom for novel algorithms. As demonstrations, we propose supervised and self-supervised implementation of T Z: M), and em- pirically show that the corresponding optimization algorithms exhibit remarkable generalization across broad spectrum of RL benchmarks, context shift scenarios. data qualities and deep learning architectures. This work lays the information theoretic foundation for COMRL methods, leading better understanding of task representation learning in the context of reinforcement learning. Given its generality, envision our framework promising offline pre-training paradigm of foundation models for decision making
Introduction
The ability to swiftly learn and generalize to new tasks is hallmark of human intelligence. In pursuit of this high-level artificial intelligence (AI), the paradigm of meta-reinforeemment learning (RL) proposes to train AI agents in trial-and-error manner by interacting with multiple external environments. order to quickly adapt 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 the new data. This meta-RL mechanism has been adopted in many applications such as games 2), robotics
3 and drug discovery (5).
However, for data collection, classical meta-RL usually requires enormous online explorations of the environments [6. [7, which is impractical in many safety- criticall scenarios such as healthcare (8) autonomous driving and robotic manipulation [10]. As a remedy, offline RL enables agents
* "qqual contribution. This work primarily done when Hai Zhang worked as intern at Zhejiang Lab. i Conresoonding Author
'Sorre code: https://gihub..cm/hetraaaya htt.:/eeraaaatteto/o//UCCCCO
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

to learn from logged experience only, thereby cireumventing risky or costly online interactions. Recently, offline meta-Re RL (OMRL) [1214 has emerged as a novel paradigm significantly extend the applicability of RI. by "killing two birds in one stone": it builds powerful agents that can quickly learn and adapt by meta-learning, while leveraging offline RI. mechanism to ensure a secure and efficient optimization procedure. In the context of classical supervised or self-supervised learning, which is de facto offline, OMRL is reminiscent of the multi-task learning
[[3,, meta-training [6] and fine-tuning [16188 of pre-trained large models. We envision it as a cornerstone of RL foundation models 19 in the future.
Along the line of OMRL research, popular paradigm that seeks optimal meta-policy conditioning on the context of Markov Decision Processes (MDPs). Intuitively. the crux of COMRL lies in learning effective task representations, hence enabling the agent to react optimally and adaptively in various contexts. To this end, one of the earliest COMRL algorithms FOCAL [12] proposes to capture the structure of task representations by distance metric learning. From a geometric perspective, essentially performs clustering by repelling latent embeddings of different tasks while pulling together those from the same task, therefore ensuring consistent and distinguishable task representations.
offline meta- vett--eenforrement learning (COMRL) is a
Despite its effectiveness, FOCAL is reported to be vulnerable to context shifts [22] i.e., when testing on -iitr----dtitbutiinn (OOD) data (Fig. I 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 (2 23 To alleviate the problem. follow-up works such as CORRO (24] reformulates the task representation learning of COMRI as maximizing the mutual information T(Z; MM between the task variable M and its latent code Z. then approximates I(Z; M) by an InfoNCE (25) contrastive loss, where the positive and negative pairs are conditioned on the same state-action tuples (8, a). Inspired by CORRO, recently proposed method CSRO (23 introduces an additional mutual information term between Z and (s, 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 266 benchmarks.
Contributions In this paper, following the recent development and storyline of COMRL, we present a Unified Information Theoretic Framework of Context-Bass Offline Betta-Reinforemment Learning (UNICORN) encompassing pre-existing methods. We first prove that the objectives of FOCAL, CORRO and CSRO operate as the upper bound, lower bound of T [Z; M) and their linear interpolation respectively, which provides nontrivial theoretical unification of these methods.
Second, by the aforementioned insight and an analysis of the COMRL causal structures, we shed light on how CORRO and CSRO improve context-shift robustness compared to their predecessors by trading off causal and spurious correlations between Z and input data X.
Lastly, by examining eight related meta-RL methods (Table T concerning their objectives and imple- mentations, we highlight the potential design choices of novel algorithms offered by our framework. As examples, we investigate two instantiated algorithms, and demonstrate ex experimentally that they achieve competitive in-distribution and exceptional OOD generalization performance on a wide range of RL domains, OOD settings, data qualities and model architectures. Our framework provides a principled roadmap to novel COMRL algorithms by seeking better apprxximaiionsregulaiiatiooo of ((Z; M), )), as well new implementations to further combat context shift.
supervised and the other self- sulf-supervised,
2 Method
2.1
Preliminaries, Problem Statement and Related Work
We consider MDP modeled as M = (S, Po. 7, H) with state space S, action space A. transition function T(s" s, a),, bounded reward function R(s. a), initial state distribution (s(s). discount factor 7 E (0, 1) and H the horizon. The goal is to find a policy π : S + A to maximize the expected return. Starting from the initial state, for each time step, the agent performs an action sampled from x, then the environment updates the state with T and returns a reward with R. We denote the marginal state distribution at time as p,, (s). The V-function and Q-function are given by

---

UNcCm RRARSOR
40
40
40
last corios
20
20
20
Y-axis
Y-axis
Yaxis
0
-20
-20
-20
-40
Lacting coctec
-40
-40
50 -50
-25
25
50
X-axis
X-axis
X-axis
Figure I: Context shift of COMRL in Ant-Dir. Left: Given a task Al' specified by goal direction (dashed line), the RL agent is trained data generated by variety of behavior policies trained on the same task M (red). At test time, however, the context might be collected by behavior policies trained on differrnnt tasks {M*} (blue), causing context shift of OOD behavior policies (Section 3.3). Middle: Against OOD context, UNICORN (red) is more robust than baselines such as FOCAL (green) in terms of navigating the Ant robot towards the right direction. Right: Besides behavior policy, the task distribution (e.g., goal positions in Ant) can induce significant context shift (Section 4.2) which is also challenging scenario for COMRL models to generalize.
H-1
V+(() = Σ 1
? E (
(1)
Q a) = R(s,a) a)
[Va (s))).
In this paper, we restrict scenarios where tasks share the state and action space and can only be distinguished via reward and transition functions. Each OMRL task is defined as an instance of MDP: M' (S, R', Po: H) M, where M is the set of all possible MDPs to be considered. For each task index i, the offline dataset X {(s),, aj, rj, "}} collected by behavior policy T... For meta- given trajectory segment ci... {(s), a), rj. "))== of length n as the context M COMRL employs a context encoder ( ci..) to obtain the latent representation z of task M'. If contains sufficient information about the task identity, COMRL can be treated as special case of RI. on partially-observed MDP [27]. where is interpreted as a faithful representation of the unobserved state. By conditioning on 2, the learning of universal policy <- z) and value function V- 2) (88 become regular RL, and optimality can be attained by Bellman updates with theoretical guarantees. this end, FOCAL [12] proposes to decouple COMRL into the upstream task representation learning and downstream offline RL. For the former, which can be treated independently, FOCAL employs metric learning to achieve effective clustering of task embeddings:
CPOCAL = min E., (1(1= sill= 2112 + 1( 2 +
(3)
+e
However, it is empirically shown that FOCAL struggles to generalize in the presence of context shift (22). This problem has been formulated in several carliest studies of OMRL. Li et al. (29) observed that when there are large divergence of the state-action distributions among the offline datasets, due to shorteut learning [30], the task encoder may learn ignore the causal information like rewards and spuriously correlate primarily state-action pairs to the task identity, leading to poor generalization performance. Dorfman al. [14] concurrently identified related problem which they termed MDP ambiguity in the context of Bayesian offline RL. An example is illustrated in Fig. T The problem is further exacerbated in the fully offline setting, the testing distribution is fixed and cannot be augmented by online exploration, allowing no theoretical guarantee for the context shift.
To mitigate context shift, subsequent algorithm CORRO (24) proposes to optimize lower bound of [ (Z; M) in the form an InfoNCE [25) contrastive loss
Ma, 2)
LOORRO = min Ez, log EM h((x,2))
(4)
3

---

where (s, as,,,,,,) s'), (0(2/=)), h(x, 2)
M(lam)
P(z] P(z) and x² (s, r,,,,,"" s'*)
transition tuple generated for task M E M conditioned on the same (s, a). To further combat the
spurious correlation between the task representation and behavior-policy-neee
state-actions, recently proposed method CSRO (3 introduces a CLUB upper bound [1 of the mutual information
between 2 and (s, a., to regularize the FOCAL objective:
LCSRO min (CrOCal + ALCLUB)
(5)
@ LCLUB IE; [log (2i a.) Ey [log u0 (z)s8, a.)]]
(6)
In the next section, will show how these algorithms are inherently connected and how their context-shift robustness gets improved incrementally.
2.2 A Unified Information Theoretic Framework
We start with a formal definition of task representation learning in COMRL:
Definition 2.1 (Task Representation Learning). Given an input context variable X E X and its associated task/MDP random variable M E M, task representation learning in COMRL aims to find a sufficient statistics Z of X with respect to M.
In pure statistical terms, Definition 2.1 limliies that ideal representation Z is mapping X that captures the mutual information I M)) We therefore define the following dependency structures in terms of directed graphical models:
Definition 2.2 (Causal Decomposition). The dependency graphs of COMRL are given by Fig. 2
where X and X, are the behavior-related (8, -CCcopponent and task-related (s', -componntttt of
the context X, with X (X,, X.).
For the first graph, M M X Markov chain, which satisfies I(Z; X 0. To interpret the second graph, given an MDP
forms
M
X
M ~ M, the state-action component of X is primarily captured by the behavior policy πg: Tg. The only exception is when tasks differ in initial state distribution Po transition dynamies T, in which case the state variable S also depends on M. We there- fore define it the behavior-reladt related component, which should be weakly causally related (dashed lines) M and Z. Moreover, when X, is given, X, is fully characterized by the transition function T (s, - s' and reward function R: (s, a) + r of M, which should be strongly
Causality Strong Causality
Mx (s),
X
M
X,
Z
Graphical Models of COMRL. causally related (solid lines) M and Z and therefore be defined as the task-related component.
Figure 2:
Mathematically, we find that rewriting X (Xe, Xs) induces a causal decomposition of the mutual
information [ZZ; X) by the chain rule 32
I( I(Z X X + I(Z;
(7)
We thereby name 1( 12 and (() the primary and lesser causality in our problem respectively. With the setup above, we present the central theorem of this paper: Theorem 2.3 (Central Theorem). Lei = denote equality up to constant, then
1 T(Z; Xx,,))
S T(Z:M) S X,, + I I(2;
T(Z;X)
primary causality
primary + lesser causality
holds up to a constant, where
1. LEOCAL = X).

---

2. LCORRO- = I(2;
3. LCSRO 2 - -(1 (22
Proof. See Appendix B
Theorem 2.3 reveals several key observations. Firstly, the FOCAL and CORRO objectives operate as an upper and lower bound of I ZZ; M) respectively. Since one would like maximize [ M) according to Definition CORRO, which maximizes the lower bound / (2; X, Xs), can effectively optimize I [ M) with theoretical assurance. However, FOCAL which maximizes the upper bound [[Z; Xx, Xt) provides guaraniee for T(Z; M). By Eq. maximizing the FOCAL objective may instead significantly elevate the lesser causality / [ X which is undesirable since it contains spurious correlation between the task representation Z and behavior policy xg. This explains why FOCAL is less robust to context shift compared to CORRO.
Secondly, CSRO as the latest COMRL algorithm among the three, inherently optimizes a linear combination of the FOCAL and CORRO objectives. In the becomes convex interpolation of the upper bound I(Z; X) and the lower bound I(Z; X; (x) of I[Z; M), which in essence, enforces trade-off between the causal (Z with Po) and spurious (Z with TT)) correlation contained in I iZ; X,). This accounts for the improved performance of CSRO compared to FOCAL and CORRO.
regime, the CSRO objective
2.3 Instantiations of UNICORN
By providing a unified view of pre-existing COMRL algorithms, Theorem 2.3 opens avenues for novel algorithmic implementations by seeking alternative approximations of the true objective I(Z; M). To demonstrate the impact of our proposed UNICORN framework, we discuss two instantiations as follows:
Supervised UNICORN I(Z can be re-expressed as
T(2; M) H(M) - H(M = - (
M= log [[ EE [M =)]
(8)
where (-) is entropy. Since in practice, each z² of sample x: is collected within specific task M minimizing the parameterized entropy Ho Z z²) is equivalent to finding optimal function Po (M]22 which correctly assigns the ground- -ruth label M to z². i.e., optimizing Pe (MM [M)2) towards delta function 6(M M²) for continuous M or an indicator function (M MM) for discrete M. This implies that
arg min Ho ( =2) = arg max log 1*)
0
Suppose a total of nM training tasks J are drawn from the task distribution p(M) with the task label M given for meta-training, Under this supervised scenario, by substituting Eq. 9 into g we have
arg max ((; = arg max E,EMs [8(M -M') log po (M'
nM
arg max Eg
Σ
1(M M) log po (z)
(10)
1=1
which is precisely the negative cross-entropy loss H(M, P(M for way-way classification with feature and classifier pg- We therefore define the objective of supervised UNICORN as
LUNICORN-SUP H(M, P(M
RM
Ex (x()))
Σ
1( Mod'=M)logg lo (M
(11)
Note that LUNICORN convex and operates as finite-sample approximation of -T(Z; M), for which we derive the following bound:
5

---

Theorem 2.4 (Concentration bound for supervised UNICORN), Denote by I M) the empirical estimate of [Z M) by nM tasks, [Z; M) the expectation, then with probability at least 1 6,
Var (ar(H(N) (2r(H(2|M))
i) M(
(12)
Proof. See Appendix B
Theorem 2.4 bounds the finite-sample estima- tion error of the empirical risk [(2; ; M) with nar task instances drawn from the real task dis- tribution p pM). The supervised UNICORN has the merit of directly estimating and optimizing the real objective [ ; M). which requires ex- plicit knowledge of the task label M and sub- stantial amount of task instances according to Theorem 2.4 For better trade-off of computa- tion and performance, we choose to sample 20 training tasks for all RL environments our experiments.
training
Xx
doocdlr
Pe(x; X.X)
gradient
ropkay
eeoade
Connect
9ala e)
2
Guía m.a)
Purts
Iasting tasks
Self-supervised UNICORN In practice, offline RI datasets may often be collected with limited knowledge of the task specifications or labels. In this scenario, previous works implement self- supervised learning 33 obtain effective rep- resentation Z, such as the contrastive- based FOCAL/CORRO to optimize I(Z; ( (//IIIIIZZ X, Xx) respectively; or generative approaches like VariBAD /34/BOReL [14] reconstruct the trajectories X by variational inference, which is equivalent to maximizing [X;; ; Z, X;). By Theorem 2.3 these methods optimize relatively loose upper/lower bound of (ZZ; M), which can be improved by convex combination of the two bounds:
Meta-learning procedure of The supervised variant UNICORN-SUP simply replaces the decoder by classifier po( p)(M]z) and optimize a cross-entropy loss instead of frcoon and LPOCAL
Figure UNICORN-SS.
I(; al (1-
(13)
where 0 < a < 1 is hyperparameter. Implementing each term in Eq. allows ample design choices, such as the contrastive losses Eqs. (3. 4 and (6) or autoregressive generation via Decision Transformer [5 36 or RNN 34 For demonstration, in this paper we employ a contrastive objective
fFOCAL, as Eq. for estimating [Z: X) while approximate T (ZZ X, Xx) by reconstruction. By the chain rule:
T(Z; X, = [(X;; Z, X.) 1(X,; Xt) = I(X;; 2, X,),
(14)
since T(X;; Xτ) is constant when X, and X, are drawn from a fixed distribution in offline RI. Moreover, by definition of mutual information:
p(Zz|zz log p(zxe)
I I(X;; X.) Eu.
Eng ab. logp(xzzl 2b,]
E 2e - (x) (og
(15)
which induces a generative objective Crecon = - ((x: Z, by reconstructing X, with a decoder
network po(-12, 2τ) conditioning Z and X,. As result, the proposed unsupervised UNICORN
objective can be rescaled as Eq. 16
LUNICORN-SS := Lrecon LEOCAL
(16)
The influence of the dyper-paramette - is shown in Appendix C.4

---

Table 1: Comparison between UNICORN instantiations and related existing contextual meta-RL methods. For clarity, 'Representation Learning Objective" only lists the loss functions of Z that are independent of the downstream RL tasks. Note that (Z Xx) (Xx: Z, X6) holds only for offline RL.
Method
Setting
Representation Learning Objective
Implementation
Context X
UNICORN-SUP
Offline
I(Z; M)
Predictive
Transition
UNICORN-SS
Offline
aI(2:X) (1 (XX;; X.)
Contrastive+Generatt
Transition
FOCAL 22 22
Offline
[Z; X
Contrastive
Transition
CORRO 24
Offline
I(2: (/,:X,)X,)
Contrastive
Transition
CSRO PX
Offline
(1- X) (2; ( )
AZ;
Contrastive
Transition
GENTLE
Offline
1(x;: Z, X²))
Generative
Transition
BOReL 14
Offline
I(X;; 2, DKL (go (Z X [Po Z))
Generative
Trajectory
VariBAD
34
Online
1(x;; (,), DxL (g X
Pol (2))
Generative
Trajectory
PEARL
Online
(gx Z X |pol
N/A
Transition
ContraBAR
42
Offline&Online
1(2 X;; A)
Contrastive
Trajectory
We illustrate our learning procedure in Fig. with pseudo-code in Algorithms ! and 2
A holistic comparison of our proposed algorithms with related contextual meta-RL methods is shown in Table [] The extra KL divergence in methods like VariBAD and PEARL can be interpreted as the result of variational approximation information between and X, which we found unnecessary our offline setting (see ablation in @) Behavior regularized actor critic (39 is employed to tackle the bootstrapping error (40) for downstream offline RL implementation.
information bottleneck [7 38) that constrains the mutual
Experiments
For brevity, we name our proposed supervised and self- superisde algorithms UNICORN-SUP and UNICORN-SS respectively. To evaluate their effectiveness, our main experiments are organized to address three primary inquiries regarding the core advantages of UNICORN: (I) How does UNICORN perform on in-distribution tasks?
(2) How well can UNICORN generalize to data collected by OOD behavior policies? (3) Can UNICORN outperform consistently across datasets of different qualities? The performance is measured by the average return across 20 testing tasks randomly sampled for each environment.
3.1 Experimental Setup
Data Collection Following FOCAL, for each task, we train SAC 43 agent from scratch and collect its replay buffer as the offline dataset. Hence our default data represent mixed distribution of behavior policies, ranging from random to expert level. Policy checkpoints are also kept for creating
various new datasets in Sections 333 and 3.4
We adopt MuloCo 26 and MetaWorld 3 benchmarks evaluate methods, which involve six robotic locomotion environments that require adaptation across reward functions (HalfCheetah- Dir, HalfCheetah- Vel,, Ant-Dir, Reach) across dynamies (Hopper-Param, Walker-Param) We compare UNICORN with six competitive OMRL algorithms, which are categorized as context-based, gradient-based and transformer-based methods:
Table 2: Average testing returns of UNICORN against baselines on datasets collected by IID and OOD behavior policies. Results are averaged by 6 random seeds. The best is bolded and the second best is underlined.
HalfCheetah-Dr Dir
HalECheetab. Vel
Ant Dic
Hopper-Paran Param
Walkeer Param
Reach
Algorithm
UNICORN-SS
1307:26
1296:24
-22a1
.9d55
267+14
236e18
31646
304::11
419244
407146
2775:241
20042183
UNICORN-S5 SUP
1296-20
1100476
-2513
-9145
2501-1
239:16
312+1
302+12
322=28
3124.99
26811111
2611:110
CSRO
11801228
1581253
-2811
-10215
276119
233112
31016
301410
310158
279165
2720u235
2801:182
CORRO
345s146
120s12
272+13
2468:175
2322x327
POCAL
1186+272
861+253
173+24
297413
308-98
286+91
2474-256
Supervised
9624356
782+429
24+1
-04ee
238+39
202+38
306+10
29468
256e60
210+28
2489=248
2783+205
MACAW
1155+10
45046
-6642
-18+1
26+3
0-0
21846
205+2
141#9
13065
2431=157
1725+79
Prompt-DTT
1176640
2549
-118466
-24+211
0=0
23445
20265
185e9
156+17
2165685
189++111

---

Ant: Dir
HolrCheetah
HolfCheetah-ve
-123
6.0
400
155
1000
13K
179
129K
100
1996
2006
Hopper-Param
Walker Pallier-PPram
Reach
30.0
0.4
2.2
123K
1230
1756 200
04
205
50K
732
500%
122K
100K
173V
2006
UNICORN-SUP
UNICORN-SS
Supervised
Prompt-DT
Figure 4: Testing returns of UNICORN against baselines 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.
FOCAL, CORRO, CSRO are context-based methods that can be seen as special cases or approxi- mations of UNICORN, see details in Section 2
Supervised is context-based method, directly using actor-critic loss to train the policy/value networks and task encoder end-to-end. We find it competitive baseline across all benchmarks.
MACAW [3] is gradient-batt based method, using supervised advantage-weighter regression for both the inner and outer loops of meta-training.
Prompt-DT is a transformer-bow baasdd method, taking context as the prompt of a Decision Trans- solve OMRL conditional sequence modeling problem.
former (DT) 35 36
Our evaluations are performed
few-shot manner. For gradient- methods, we first update the meta-policy with a batch of testing context and then evaluate the agent. For transformer-bee methods, we utilize trajectory segments as the prompt to condition the auto-regressive rollout.
3.2 Few-Shot Generalization to In-Distribution Data
For each environment, the meta-policy is trained 20 random tasks. To test IID generalization, we sample one trajectory for each testing task as the context, which enables few-shot adaptation of the agent via updated task representation F illustrraees the learning curves of UNICORN vs. all baselines in terms of testing returns, which correspond to the IID entries of Table 2 The results demonstrate that UNICORN variants, especially UNICORN-SS, consistently achieve SOTA performance across all benchmarks. The observation that the asymptotic performance of CSRO comparable to UNICORN is expected since it also optimizes linear combination of the lower and upper bound of (Z;M) by Theorem 2.3 with a different implementation choice. Notably, UNICORN exhibits remarkable stability, especially on HalfCheetah- and Palker-Paramm, where the other context-basee process. Although the gradient-basee MACAW and the transformer-basee Prompt-DT show faster convergence, their asymptotic performance leaves much to be desired. Moreover, the average training time for MACAW about three times longer than the context-bant based methods.
methods suffer a significant performance decline during the later training
3.3 Few-Shot Generalization to Out-of-Distribution Behavior Policies
Denote by {T's T==1=mar,t=0- the checkpoints of behavior policies for all tasks logged during data collection, where i refers to the task index and represents training iteration. To evaluate OOD generalization, for each test task M we use all behavior policies in TTT's 6.t )i to collect the OOD contexts and the conditioned rollouts. The OOD testing performance is measured by averaging returns across testing tasks and behavior policies, as shown in the OOD entries of Table While all methods suffer notable decline when facing OOD context, UNICORN maintains the strongest

---

performance by significant margin. Another observation is that the context- based methods are in general significantly more robust compared to the gradient- based MACAW and transformer-basee Prompt-DTT. This also justifies the storyline of COMRI development from 2021 to date
FOCAL + CORRO + (2021)
CSRO + UNICORN (2023)
(2022)
(2024)
as roadmap for pursuing more robust and generalizable task representation learning for COMRL.
3.4 Influence of Data Quality
To test whether the UNICORN framework can be applied to different data distributions, we collect three types of datasets whose size equal the mixed dataset used in Section 3.2 and 3.3) random, medium and expert, which are characterized by the quality of the behavior policies.
As shown in Table UNICORN demonstrates SOTA performance all types of datasets. Notably, despite CSRO having better IID results than UNICORN on Ant-Dir under the mixed distribution, UNICORN surpasses CSRO by large margin under the narrow distribution (medium and expert) and random distribution.
We observe that CORRO fails all three datasets, which is likely attributable to the re- liance of CORRO
Table 3: UNICORN vs. baselines Ant-Dir datasets of various qualities. Each result is aver- aged across 6 random seeds. The best is bolded and the second best is underlined.
the negative sample gen- erator. On one hand, under narrow data dis- tribution, there might be little overlap among the high-density state-action regions of datasets from different tasks. On the other hand, sam- ples collected by random policies are often in- distinguishable across tasks (22). Both factors make it more challenging to train a robust gener- ator for producing high quality negative samples conditioning on specific state-action pairs. As for Prompt-DT, its performance improves sig- nificantly with increased data quality, which is expected since the decision transformer adopts a behavior-eloning-lne like supervised learning style.
Random
Medicam
Expert
Algorithm
HD
0OD
IID
00D
IID
0OD
UNICORN-SS
81#18
6246
220=23
243+10
279:10
262+13
UNICORN-SUP
75a15
5045
140=11
126+32
247=15
229:19
CSRD
2+3
0+1
166-10
198+17
202-45
CORRO
1+1
[H]
8-5
-742
-++10
-14=9
ROCAL
67626
44:10
171:84
187186
229:44
246120
Suparvised
6566
47+12
149-50
110+80
249H33
215+60
MACAW
040
28=2
S8:43
12l
Prompt-DT
264
78a15
4 Discussion
This section presents more empirical evidence on the applicability of the UNICORN framework.
4.1 Is UNICORN Model-Agnostic?
As UNICORN tackles task representation learning from a general information theoretic perspective, it is in principle model-agnon aggostic. Hence a natural idea to transfer UNICORN other model architectures like DT, which we envision to be promising backbone for large-scale training of RL. foundation models.
Table 4: DT implementation of COMRL on
HalfCheetah-Dir and Hopper-Param. Each result is averaged by 6 random seeds.
Hal:Chectab--De
Hopper-Param
Algoritlun
IID
OOD
II)
0OD
UNICORN-SS
1307=26
1296x24
31616
304a11
We employ a straightforward implementation by em- bedding the task representation vector the first token in sequence to prompt the learning and gen- eration of DT, akin to in-context learning for large
UNICORN-SS-DI
1233=10
1186143
30424
291
UNICORN-SUP SUP-N-SUP-DT DT
1227=21
1065457
30846
29712
POCAL-DT
1209=33
652436
29344
28485
Prompt-D
1177:40
-2569
234+5
20355
Instead of the non-trane non-trainable prompt used in the enforce task representation learning of the prompt by optimizing LuNICORN-Ss. LENICORN.SUP and LFOCAL. We name these variants UNICORN-SS-DT. UNICORN-SUPA and FOCAL-DT respectively.
language models 45 and in-context RL [46] [49]. original Prompt-DT,
As shown in Table 4 and Fig. 7 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 UNICORN- SS 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 [50] and DT 51.
We leave this verification to future work.
4.2 Can UNICORN be Exploited for Model-Based Paradigms?
Ant-Dir
With decoder po(a, 2, zb)) as world model [52], UNICORN-SS can potentially enable a model- -ased paradigm by generat- ing data using customized latent 2, which can be especially useful for generalization to OOD tasks. We implement this idea on Ant-Dir by sorting the tasks accord- ing to the polar angle of their goal direc- tion (0 E [0, 27) and construct task-level OOD by taking 20 consecutive tasks in the middle for training and the rest for testing (Fig. [ Gaussian noise e is ap- plied to the real task representation z and (s, a, x + e) is then fed to the decoder generate (s'; r). The imaginary rollouts
30
-30
CaRO
25K
125K
150K
175K
200K
steps
Testing returns for OOD tasks. The learning curves are averaged over random seeds.
Figure 5:
((1,, at, 8 rt)) ==) are used during the training of RI. agent. We adopt the conventional ensemble technique in model-based RL
5 54
to stabilize the training process.
As shown in Fig. 5 SANICORN-SS-enbbll performance (indicating positive few-shot transfer) in this extremely challenging scenario where context shift is induced by OOD tasks.
model-based RL is the only method achieve positive
5 Conclusion & Limitation
In this paper, we unify three signature context-based offline meta-RL algorithms, FOCAL, CORRO and CSRO (potentially a lot more) into a single framework from an information theoretic rep- resentation learning perspective. Our theory offers valuable insight how these methods are inherently connected and incrementally evolved in pursuit of better generalization against context shift. Based on the proposed framework, we instantiate and UNICORN-SS which are 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 tolioie- COMRL. Moreover, we expect our framework to incorporate almost all existing COMRL methods that tackle task representation learning. Since according to Definition 2. as long as the method tries to solve COMRL by learning sufficient statistic Z of X w.r.t. M, will eventually come down optimizing an ieformation-thhoree objective equivalent to I(Z. M), a lower/upper bound approximation like the ones introduced by Theorem 2.3] up to some regularizations or constraints. With such mathematical generalization as well as the empirical evidence of UNICORN being model-agnostic (see Section 4. 1). we envision it as a nascent offline pre-training paradigm of foundation models for decision making 21
novel algorithms called UNICORN-SUP
Limitation Since our framework assumes a decoupling of task representation learning and offline policy optimization, it does not directly elucidate how high-quality representations are indicative of higher downstream RL performance. Another limitation would be the scale of the experiments. Our offline datasets cover most 40 tasks and 450k transition tuples, which might be the reason why UNICORN-SUDE is inferior to UNICORN-SS (n too small to reduce the approximation error in Theorem 2.4p However, we expect our conclusion extrapolate to larger task sets, dataset sizes and model parameters, for which the DT variants may demonstrate better scaling. Lastly, in principle, the information theoretic formalization in this work should be directly applicable to online RL. However, many of our key derivations rely on the static assumption of M and X (see Appendix B for details), which are apparently violated the online scenario. Extending our framework to online RI. interesting and nontrivial, which we leave for future work.
10

---

Acknowledgments and Disclosure of Funding
The authors would like to thank the anonymous reviewers for their insightful comments and sug- gestions. The work described in this paper supported partially by grant from the Research Grants Council of the Hong Kong Special Administrative Region, China (Project Reference Number: T45-401/22-N), by a grant from the Hong Kong Innovation and Technology Fund (Project Reference Number: MHP/086/21) and by the National Key Research and Development Program China (No. 2020YFA071 1402).
References
[1] Mandi Zhao, Pieter Abbeel, and Stephen James. On the effectiveness fine-tuning versus meta- reinforcement learning. Advances Neural Information Processing Systems, 35:26519-2655 2022.
[2] Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. A survey of meta-reinforcement learning. arXiv preprint arXiv:2301.08028, 2023.
[3] 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. In Conference robot learning, pages 1094-1100, PMLR, 2020.
[4] Ashish Kumar, Zipeng Fu, Deepak Pathak, and Jitendra Malik. Rma: Rapid motor adaptation for legged robots. arXiv preprint arXiv:2 040344 2021.
[5] Leo Feng, Padideh Nouri, Aneri Muni, Yoshua Bengio, and Pierre-Lue Bacon. Designing biological sequences via meta-reinforcement learning and bayesian optimization. arXiv preprint arXiv:2209.06259, 2022.
[6] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. International conference on machine learning, pages 1126-1135. PMLR, 2017.
[7] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off off-policy meta-reinforcement learning via probabilistic context variables. In International conference on machine learning, pages 5331-5340. PMLR, 2019.
[8] Omer Gottesman, Fredrik Johansson, Matthieu Komorowski, Aldo Faisal, David Sontag, Finale Doshi- Velez, and Leo Anthony Celi. Guidelines for reinforcement learning in healthcare. Nature medicine, 35(1):16-18 2019.
[9] Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey, IEEE Transactions on Intelligent Transportation Systems, 23(6):4909 4926, 2021.
[10] Di Zhang, Bowen Lv, Hai Zhang, Feifan Yang. Junqiao Zhao, Hang Yu, Chang Huang, Hongtu Zhou, Chen Ye, and Changjun Jiang. Focus on what matters: Separated models for visual-based rl generalization. arXiv preprint arXiv:2410. 10834, 2024.
[11] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.016433 2020.
[12] Lanqing Li, Rui Yang, and Dijun Luo. Focal: Efficient fully-offline meta-reinforcement learning via distance metric learning and behavior regularization. In International Conference on Learning Representations, 2021.
[13] Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning, pages 7780-7791. PMLR, 2021.

---

[14] Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning identifiability challenges and effective data collection strategies. Advances in Neural Information Processing Systems, 34:4607 461886 2021.
[15] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-555 2020.
[16] Edward J Hu, Phillip Wallis. Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2021.
[17] Xiang Lisa Li and Percy 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 Long 2021.
Papers), pages
[18] Brian Lester, Rami AI-Rfou, and Noah Constant. The power of scale for parameter-efficien prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 3045-3059, 2021.
[19] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the 2021.
opportunities and risks of foundation models. arXiv preprint
[20] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel Barth-maron, Mai Giménez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. generalist agent. Transactions on Machine Learning Research, 2022.
[21] Sherry Yang. Ofir Nachum, Yilun Du Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foun- dation models for decision making: Problems, methods, and opportunities. arXiv preprint vrXiv:2303.04129 2023.
[22] Lanqing Li, Yuanhao Huang, Mingzhe Chen, Siteng Luo, Dijun Luo, and Junzhou Huang. Provably improved context-based offline meta-rl with attention and contrastive learning. arXiv preprint arXiv:2 2021.
[23] Yunkai Gao, Rui Zhang, Jiaming Guo, Fan Wu, Qi Yi, Shaohui Peng, Siming Lan, Ruizhi Chen, Zidong Du, Xing Hu, et al. Context shift reduction for offline meta-reinforcement learning. In sevethh Conference on Neural Information Processing Systems, 2023.
[24] Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reinforcement learning via contrastive learning. In International Conference on Machine Leurning, pages 25747-25759. PMLR, 2022.
[25] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 1807 00344, 2018
[26] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 JEEE/RSJ international conference on intelligent robots and systems, pages 5026-5033. IEEF, 2012.
[27] Leslie Pack Kaelbling, Michael L. Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-134, 1998.
[28] Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function ap- proximators. In International conference machine learning, pages 1312-1320. PMLR, 2015.
[29] Jiachen Li, Quan Vuong, Shuang Liu, Minghua Liu, Kamil Ciosek, Henrik Christensen, and Hao Su. Multi-task batch reinforcement learning with metric learning. Advances in Neural Information Processing Systems, 33:6197-6210, 2020.
12

---

[30] Robert Geirhos, Jern-Kerrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix Wichmann. Shortcut learning deep neural networks. Nature Machine Intelligence, 2(1 2(11):665--)) ):665-673, 2020.
[31] Pengyu Cheng, Weituo Hao, Shuyang Dai, Jiachang Liu, Zhe Gan, and Lawrence Carin. Club: contrastive log-ratio upper bound of mutual information. In International conference on machine learning, pages 1779-1788. PMLR, 2020.
[32] Raymond W Yeung. Information theory and network coding. Springer Science & Business Media, 2008.
[33] Xiao Liu, Fanjin Zhang, Zhenyu Hou, Li Mian, Zhaoyu Wang. Jing Zhang, and Jie Tang. Self-supervised learning: Generative contrastive. IEEE transactions knowledge and data engineering, 35(1):857-876, 2021.
[34] Luisa Zintgraf. Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal. Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep via meta- learning. In International Conference on Learning Representations, 2020.
[35] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 34: 14:15504-55097
2021.
[36] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273- 1286, 2021.
[37] Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 ieee information theory workshop (itw), pages 1-5. IEEE, 2015.
[38] Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. In International Conference on Learning Representations, 2017.
[39] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint 191.v:111..11999. 2019.
[40] Aviral Kumar, Justin Fu, Matthew Soh. George Tucker, and Sergey Levine. Stabilizing off- policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
[41] Renzhe Zhou, Chen-Xist Gao, Zongzhang Zhang, and Yang Yu. Generalizable task representa- tion learning for offline meef-reiiforeemeett learning with data limitations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 17132-17140, 2024.
[42] Era Choshen and Aviv Tamar. Contrabar: Contrastive bayes-adaptive deep rl. In International Conference Machine Learning. 2023.
[43] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey 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
[44] Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In international conference on machine learning, pages 24631-2464 PMLR, 2022.
[45] Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng. Jingyuan Ma, Rui Li. Heming Xia, Jingjing Xu, Zhiyong Wu, Baobao Chang, et al. A survey on in-context learning. In Proceedings of the 2024 Conference on Empirical Methods Natural Language Processing, pages 1107-1128, 2024.
[46] Michael Laskin, Luyu Wang, Junbyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse, Steven Stenberg Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. In The Eleventh International Conference on Learning Representations, 2023.
13

---

[47] Licong Lin, Yu Bai, and Song Mei. Transformers decision makers: Provable in- in-oontext reinforcement learning via supervised pretraining, In The Twelfth International Conference on Learning Representations, 2024.
[48] Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brumskill. Supervised pretraining can learn in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
[49] Zhi Wang, Li Zhang, Wenhao Wu, Yuanheng Zhu, Dongbin Zhao, and Chunlin Chen. Meta-DT: Offline meta-RL as conditional sequence modeling with world model disentanglement. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
[50] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836), 2020.
[51] Kuang- Lee, Ofr Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadarrama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. Advances Neural Information Processing Systems, 35:27921-27936, 2022.
[52]
David Ha and Jürgen Schmidhuber. World models. arXiv preprint /83. 00122 2018..
[53] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33: 1111412--1442, 2020.
[54] Hai Zhang, Hang Yu, Junqiao Zhao, Di Zhang, Chang Huang, Hongtu Zhou, Xiao Zhang, and Chen 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.
[55] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 911)) 2008.
14

---

A Pseudo-Code
Algorithm UNICORN Meta-training
Ji=l Require: Offline Datasets X { training tasks M {M') Niii initialized learned policy
Two Q function Qv, context encoder 9Φ. classifier/decoder pe, hyper-parameter a, task batch
size Nes, learning rates B1, A3,, As, B4
1: while not done do
2: Sample task batch {M²) M and the corresponding replay buffer R {Xi) Ni X
3:
for step in each iter do
4:
Train the context encoder and decoder
5:
Sample context {ct 1 R
6:
Obtain task embeddings {ar ~ (=(2)c²)) Na
7:
Estimate LUNICORN (CUNICORN-SUP or LUNICORN-SS)
8:
Φ Φ- 81 V. &
9:
-- 32 VeCUIICONN
10:
// Train the actor and critic
11:
Detach the task embeddings {2/) 1
Nob
Sample training data {dl Nox R Estimate Lactors Leritic
13:
14:
w x- w B3 V fontoor
15:
1/2 + Ba V Ceritic
16:
end for
end while
Algorithm 2 UNICORN Meta-testing
Require: Offline Datasets X {X
Nus testing tasks M = {M*) Ii=l
Mus
trained policy Tu and
context encoder 4
1: for each M' do
2: Sample context e¹ ~ X2
3: Obtain task embedding 2t (0(2)c)
4: Rollout policy Tw (als, 2²) for evaluation
5: end for
B Proofs
B.1 Proof of Theorem 2.3
We first prove the following lemma:
Lemma B.1. For COMRI, T(Z: M) 2 I(Z; M Mt);; or equivalently, T(Z: M; X 2 0, both up constant.
Proof.
XX, =0
= +[ZZ M) - [(2; X,) = [(M; X,, X,) 7(M; X,, X,12) 20
>IMM; Xx) + I(M; X() I(M: X, Z) [(M; X, X,, 2
(MM;; X²) - [IM; 2 -I(M; X, X1,) + [(M; X, Xb. Z > const.
const
20
Now, since (Z;M) (( M ((M; I(M; Xx we have 1(Z;
I(Z; M ((Z; M: X,) b) 2 const = O.
With Lemma B.I proceed to prove Theorem 2.3 follows:
15

---

Proof. I X; CCORRO & [(; MM) is shown in the Theorem 4.1. of CORRO [24] and GENTLE 41 with strong assumptions. We hereby present a proof with assumptions. Given
T(Z; M; Xx) 2 0 (Lemma B.1). we have
M I(Z (6) + M
I(Z; (
T[MM; Z, X X1)
- [[M: X, 2,
-I(Z Xb)
+1(M: X, Xx) I(M; X, Xb)
Ob M-X-2
-I(M: X,)
[MM X, X;)
I(M; X, X6) H(X, Z, X.) H(X, 2,X;)
2 [MM; (o)) H(X, Z. Xb)
I(M; X
(()) +I(X Z.)
T(X;; Z, X6)
[[Z; Xt) XI)
cONSE
[(Z X,
as desired. The third and fourth lines are obtained by direct application of the chain rule of mutual information All mutual information terms without Z are held constant in the fully offline scenario.
IZ M) < X) is a direct consequence of the Data Processing Inequality [32 and the Markov chain M + X - Z in Definition 2.2 We now prove the three claims regarding FOCAL, CORRO and CSRO:
1. CPOCAL = -1(Z
p(z, x) p(2 (z))p(x)
1(2; X): = Ex,
log
log
+ +log([M))
p(* (2) MM|
p(z)
Eg,s
log
+ log(l-M)
p(z) p(E|2)
MM
p(= p(z)
~ Eg,z
log
+ log(lM))
x(z) p(z|=)
Σ EMEEM Exixxi
plalan lan")) p(z)
p(z(ar) p(z)
log
log(l-M)
EMSEM Exixx
p(x)(x))) p(2)
h(1, 2)
)
En.z
log
+ log((M))
EMPEM h((x),)))
The first term on RHS precisely the supervised contrastive learning objective introduced by a variant of FOCAL 22). which is equivalent to the negative distance metric learning loss CROCAL. with the effect of pushing away embeddings of different tasks while pulling together those from the same task. Therefore have
16

---

EMiEM h(20',2) )
1( 2)
((2 = Ex, log
+ const
-FFCCAL.
2. LCORRO = - X((x))
p(z, 2, xee (26)
((::X, = Ex,,
log
p(= (26 (ppp)))((( (p(ce)p(eeee (206)
log
+ log(
p(2(2b))
p(z me .66
(66
log
+ log(l
mL) X
M BMrEM
p(z .20)))) p(z|wa)
2 ,0
log
+ log((M))
p==(mb) p( (b6)
ΣM-EN EM
p .k)
p(=
p(z) ate .26) pl (9%)
log
log(I-M)
p(=| EM-EM p(= mb)
p(x (0))
log
p(2)
+ log(lM))
EM-eM
p(#12) 2b) p(s)
h(x,2)) z)
)
R log
EM-EM 1(x",,,0 )x)
+ log((M)
= -CORRO.
3. LESRO 2 (A 1) - Ax
The CLUB loss in Eqn 6 operates upper bound of (,) 31
By Eqn
5
we have
LCSRO
LFOCAL ALCLUB
2 LHOCAL + A[(2X;))
LPOCAL -((2
mutuaal information
+ X(T [(Z; x, Xx))
- ()) -
B.2
Proof of Theorem 2.4
Proof. Denote by nz aM [ |X*| the number of context samples, we have
1Z; H ((Z))
(17)
log p(z))
EM H = M))
(18)
n2
RM Σ j=l
log p(2))
H M³).
(19)
17

---

Since >> n.M, the concentration characteristic of I [2; M) is dominated by the second term. If we ignore the approximation error of the first term, by the Chebyshev's inequality, for any 0.
Var(H Z M))
Pr ( 26
(20)
nM
Or equivalently, with probability at least 1 8,
Var (Z)M)
f
(21)
nyr
Remark:
Consider practical implementation of M)) as multivariate Gaussian N(u., EMr),
which gives H (Z M) [og](27e)dmmm2 [3 probability least. 1 5,
Substituting into Eqn 21 we have with
Var ΣMr
-
i(z:M) T(Z:M) &
(22)
C Further Experiments
C.1 UNICORN-SS-O: A Label-free Version
When a = O. UNICORN-SS task representation learning reduces to the minimization of the recon- struction loss only, which becomes a label-free algorithm (i.e., it does require the knowledge of task identities/labels to optimize, as opposed to the contrastive learning in previous methods). We name this special case UNICORN-SS- which equivalent For fair comparison, we use a variant of a Bayesian OMRL method BOReL [14] that does not utilize oracle reward functions, label-free baseline. test time, we modify BOReL to use offline datasets rather than the online data collected by interacting with the environment to infer the task information As shown Table UNICORN-SS-O also competitive with BOReL.
concurrent method GENTLE [4I].
Table 5: ANICORN-SS-C compared to another label-free COMRL method, BOReL, Ant-Dir.
Ant-Dir
Algorithm
IID
OOD
UNICORN-SSS
220+16
20019
BOReL.
206+18 18
187+10
To further validate that the KL constraint D (=a())=)) /)) is unnecessary, we weight the KL constraint to UNICORN-SS-O. Table|sshoww that KL. constraint would reduce the performance under the offline setting.
C.2 Visualization of Task Embeddings
For better interpretation, we visualize the task representations by 2D projection of the embedding vectors via t-SNE [55]. For each test task, generate trajectory using each policy in { }., and infer their task representations This setting aligns with our OOD experiments in the main text.
Since the UNICORN-SS objective operates convex combination of the reconstruction loss (UNICORN-SS-O) and the FOCAL loss, we compare representations of UNICORN-SS to these
è Experiment shows that the usage of offline dataset rather than online data has little effect on the results of BOReL.
18

---

Table 6: UNICORN-SS-0 compared to other KL constraint weight variants.
Ant-Dir
Algorithm
IID
OOD
UNICORN-SS-0
220+16
200:9
KL-0.1
215+11
189+12
KL-0,5
202+17
182+13
KT.1100
194+11
161+7
KL-5.0
162414
143+10
two extremes. As shown in Figure 6 compared to UNICORN-SS-O, UNICORN-SS 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-SS. We speculate that since FOCAL blindly separates embeddings from different tasks, it may fail to capture shared structure between similar tasks, whereas UNICORN-SS is able to do so by generative modeling via the reconstruction loss.
UNICORN-0
UNICORN
FOCAL
Figure 6: 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 mapped to rainbow colors, ranging from purple to red.
C.3 More on UNICORN-DT
We provide the learning curves of the results in Section 4.1 here. As shown in Figure UNICORN-SS: DT and UNICORN-SUU SUNN--UP-DTT demonstrate faster convergence speed compared vanilla UNICORN- SS and a higher asymptotic performance compared to vanilla Prompt-DT.
HellCheetah--ii
Horpee-Pamm
1200
3x00gg 010009-000m
evenge retum retum
800
100
0K
25K
50K
75K
100K
125K
150K
175K
200K
0K
28K
50K
70K
125K steps
100K
160K
175K
200K
INCORN-SUP-DT
INICORI-SS-DIT
LINICRNN-SS
PromphoT
POCA--TT
Figure 7: Test returns of UNICORN-SUP-DT and UNICORN-SS-DT against UNICORN-SS, Prompt- DT and FOCAL-DT benchmarks. The learning curve averaged by random seeds.
19

---

C.4 Ablation Study
Ant-Dir
Algorttmm
250
FOCAL
200
0.07
0:03
average return
150
100
50
0
0K
25K
50K
75K
100K steps
125K
150K
175K
200K
Figure 8: Different hyper-paramette settings of T-o on Ant-Dir. The learning curve is averaged by random seeds.
To illustrate the effect hyper-parameter I on the asymptotic performance, we set the following ablation study. Figure shows that as I-a increases, the performance gradually increases but when goes excessively the performance would decrease (FOCAL means = oc), which validates our proposed theory. We also find that UNICORN-SS maintain relatively stable performance range of I-a values, thus alleviating the exhaustion from parameter-tuning.
D Experimental Details
Table 7 lists the necessary syper-parameters that used to produce the experimental results.
Table 7: Configurations and hyper- parameters used in the training process.
Configurations
Ant-Dir HallCheetah-Dir
HalfCheetah- Vel
Hopper-Param
Walker-Param
Reach
daraset size
2c5
2e5
3e5
4.5e5
le5
task representation dimension
40
32
2
weight t
traininn
200k
RL. batch size
256
context training
mijectory (200 steps)
trajectory (500 steps)
learning rate
3e.4
RL network width
256
depth
cacoder width
61
128
64
20

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the main contributions are clearly laid out at the end of introduction. See Section 2 for our theory and Section for our experimental verification.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Yes, see Section 5
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model sell-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
21

---

Justification:
Yes, see for example Section 2.2 and Appendix B for complete proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Yes, see Appendix D for experimental details. Source code is provided in the Supplementary Material.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
22

---

Answer: [Yes]
Justification:
Yes, see Supplementary Material for the source code and README instruc-
tions.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cr ec/
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https i
//nips. //nips.cc/pil wwwppbblccgguiicuceccccc
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, see for example Appendix C.4 for hyperparameter study.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Yes, all figures and tables contain error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
23

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources Question:
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: See README in Supplementary Material for computing environment infor- mation.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Ethics.
We make sure to preserve anonymity and conform to the NeurIPS Code of
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
The societal/general impacts of offline meta-RI. are discussed in Section
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We only use open-source benchmarks for experiments: MuJoCo (Apache-2.0 license and Meta-Wort (MIT license).
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt
package should be provided. For popular datasets, paperswi thcode. cow/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
25

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subject.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
26


Reviews:
[{"review_id": "tnJbNJX1fM", "reviewer": "Reviewer_sHdF", "summary": "This work introduces two new algorithms for COMRL based on an information theoretic decomposition of behaviour and environment information. The primary insight is that when encoding the contextual task representation as much environment information should be maintained while as little behaviour information is maintain. The UNICORN-SS and UNICORN-SUP algorithms ultimately build on the FOCAL algorithm with the addition of either a reconstruction or classification loss. In the case of classification loss, a portion of the training network must predict which task is being performed from the inferred latent task representation and behaviour component of the context. In the reconstruction loss, a portion of the network must reconstruct the next state and reward - the task component of the context - from the latent representation and behaviour component. Experiments show that UNICORN-SS and UNICORN-SUP have more consistent performance on in-distribution samples and superior performance on out-distribution samples.", "questions": "1. I'm not certain on what is being presented in Fig 1a) with the red line. The message of the figure is clear, but if the red line is showing behaviour of an agent on the task then why would it not be more correlated with the goal direction?\n2. The last paragraph of page 4 makes the claim that using a context combination of causalities might explain why CSRO outperforms FOCAL and CORRO. It is not clear to me why including any lesser causality would improve the latent representations. Since this influenced UNICORN I feel this is  important for me to understand\n3. In Table 1, column Context  X. What does \"Transition\" vs \"Trajectory\" mean?\n4. For Equation 14. Why would switching to sampling from the encoder guarantee that the expected log probability will decrease from the line  above? Is there an assumption here or something I am missing?\n5. Figure 3. was extremely useful but even still I find that I am unsure exactly of how the context is obtained from the replay buffer. Is it correct that the latent task variable is inferred from a subset of the replay buffer with multiple steps, while the $x$ variables are a single sample which is being trained on? Perhaps annotation the middle and bottom arrows from context as you did with the top branch would be helpful.\n6. In Table 3. Surely no algorithm should work for a random policy. By definition the random policy is not goal directed and so it should be impossible to infer task from this dataset with the same state and action space?", "limitations": "Limitations are stated and clear.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "# Originality\nThe decomposition of context into behaviour and task information is interesting and allows for a clear classification of previous COMRL tasks. This more abstract grouping of algorithms is helpful and new.\n\n# Clarity\nOverall this work does an excellent job of introducing multiple algorithms and consideration clearly. The main differences between FOCAL, CORRO and CSRO are clear and intuitive from the equations presented. The notation used is intuitive and clear, and the paper is well written. Figures are clear and easily understadible.\n\n# Quality\nThe aim of the work is clear and the logic and structure of the argument is excellent. Experiments are appropriate for assessing the utility of the new algorithms and the the baselines are challenging and the appropriate choices to my knowledge.\n\n# Significance\nThe empirical results of this work are convincing for me that this indeed makes a significant contribution to the field. I appreciate the improved consistency on in-distribution performance as the improved out-distribution performance as getting less variance in RL tasks is not easy in its own right. Overall I think this is a very good paper.", "weaknesses": "My weaknesses are relatively minor.\n\n# Clarity\nLegends on Fig 1 should be larger. There were also a couple statements which were unclear but I ask about these in Question below. The statement in Theorem 2.4 was difficult to understand and I relied on Equation 11 for meaning. I think it could just be a grammar error. As a suggestion I would put the equality $I(Z;X) = I(Z; X_t | X_b) + I(Z;X_b)$ in the inequality of Theorem 2.3.\n\n# Significance\nI do think the theorems need proof sketches. This is because those theorems are the main source of new insight from this work - particularly Theorem 2.3. Due to the nature of information theory proving this theorem is likely very familiar and the construction of the COMRL problem to fit the format is the bigger advance, I see that, but it would still be helpful to get some brief intuition for how the result emerges in the main text. It would also be helpful if it was more explicitely stated why $L_{focal}$, $L_{CORRO}$ and $L_{CSRO}$ are grouped into those mutual information types. Looking at their individual equations I can see from the variables present why these types were allocated but the clearer statement on why would be useful. My last weakeness on significance is one the authors acknowledge (which I appreciate already) - due to the broad nature of the mutual information theory representation it is not clear why UNICORN outperforms the other CORRO or another algorithm which operates on the primary causality or a convex combination of causalities. I do not however see how the authors could account for this and so I make this point lightly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the constructive comments and the appreciation of our work. Our response to your concerns/questions:\n\n# Clarity\n\nThank you for your advice. We will make sure to improve the presentation of figures and theorems in the final version. \n\nTheorem 2.4 mainly bounds the error of estimating the real expectation $\\bar{I}(Z; M)$ with $n_M$ finite task instances drawn from the real task distribution $p(M)$. Please check our general response for in-depth discussion.\n\nAdding $I(Z; X) = I(Z; X_t|X_b) + I(Z; X_b)$ in Theorem 2.3 is a good idea.\n\n# Significance\n\n- proof sketch: \n\nIn the famous InfoNCE paper [1], it is shown that the contrastive loss $L_N$ as in FOCAL [2] upper-bounds the mutual information of original signal $x$ and its representation $z$: $I(Z, X) \\ge \\log (|M|) - L_N$. If we sample the positive and negative pairs in $L_N$ according to the real task distribution $p(M)$, we can turn this inequality to equality, which proves that $L_{FOCAL} \\equiv -I(Z, X)$. For CORRO, it's the same format except that all samples are conditioned on the same $x_b$, which is equivalent to the primary causality $I(Z, X_t|X_b)$. For CSRO, the additional CLUB loss in Eqn 6  provides an upper bound of the lesser causality $I(Z, X_b)$. By combining it with the FOCAL loss ($I(X, Z)$ = $I(Z; X_t|X_b)$ + $I(Z, X_b)$), the CSRO loss is effectively a linear combination of the primary and lesser causalities.\n\n- why UNICORN outperforms baselines: \n\nAs explained in our following reponse to Q2, the lesser causality $\\textit{does}$ contain some causal information about the task identity, therefore Thm 2.3 establishes the information-theoretic superiority of UNICORN over FOCAL and CORRO, since FOCAL introduces spurious correlation whereas CORRO overlooks part of the causal information. The comparison between UNICORN and CSRO is trikier, since the objective of CSRO and UNICORN-SS are mathematically equivalent. We believe the main difference lies in the implementation choice. CSRO approximates $I(Z; X_b)$ with the CLUB loss. On top of this approximation, it employs variational distribution $q_{\\theta}(z|x_b)$ to approximate the intractable real conditional distribution $q(z|x_b)$ in Eqn 6. Therefore it involves $\\textit{2 levels}$ of approximation. In contrast, both UNICORN variants involves only ${1 level}$ of approximation (finite sample approximation in Thm 2.4 for UNICORN-SUP and variational approximation in eqn 14 for UNICORN-SS). We believe this explains (at least partially) the superiority of UNICORN over CSRO. Please check our general response for in-depth discussion.\n\n# Questions\n\nQ1. \"why would the red line in Fig 1a) not be more correlated with the goal direction?\"\n\nThe red lines consist of 20 trajectories randomly picked from our mixed-quality dataset (used in sec 3.2, 3.3), which contains trajectories logged throughout the entire training cycle of a behavior policy agent. Therefore it should be a mixture of random, medium and expert-level trajectories, which is why they are not sufficiently correlated with the goal direction. \n\nQ2. \"why including any lesser causality would improve the latent representations?\"\n\nThe joint distribution of (s, a) depends on several factors: the initial state distribution $\\rho_0$, the transition function $T$ and the behavior policy $\\pi_{\\beta}$. Therefore even it's highly correlated with $\\pi_{\\beta}$, it does in principle contain $\\textit{some}$ causal information about the task identity $M$, especially for tasks that differ in transition dynamics (e.g., Walker, Hopper). Completely dropping the term might be an overkill. We briefly explained this in line 159-161 but will further clarify in the final version.\n\nQ3. \"Transition\" vs \"Trajectory\"\n\n\"Transition\" corresponds to the single-step tuple $(s_t, a_t, r_t, s_{t+1})$, whereas \"trajectory\" stands for a complete history $(s_0, a_0, r_0, s_1, \\ldots, s_t)$ of MDP up to a certain time step $t$. \n\nQ4. Last inequality of eqn 14.\n\nSince the true $q(z|x)$ is unavailable, it's a common practice in variational inference [3] to approximate it by the posterior of the encoder $q_{\\theta}(z|x)$. The inequality is caused by approximating the true probability $p(x_t|z, x_b)$ with a decoder $q(x_t|z, x_b)$\n\nProof: let $q$ denote $p_{\\theta}$ in eqn 14,\n\n$-I(X_t; Z, X_b) \\equiv -\\mathbb{E}_{x_t, x_b, z} [\\log p(x_t|z, x_b)]$\n\n$= -\\mathbb{E}_{x_t, x_b, z}[\\log \\frac{p(x_t|z, x_b)}{q(x_t|z, x_b)}]$\n\n$-\\mathbb{E}_{x_t, x_b, z}[\\log q(x_t|z, x_b)]$\n\n$= -\\mathbb{E}[KL(p(x_t|z, x_b)||q(x_t|z, x_b))] + L_{\\text{recon}} \\le L_{\\text{recon}}$\n\nQ5. \"how the context is obtained from the replay buffer\"\n\nFor context encoder: we randomly sample one trajectory $\\tau$ from the replay buffer, compute $z$ for every single-step transition $x$ in $\\tau$, and take the averaged $\\bar{z}$ as the representation.\n\nFor downstream RL training: we randomly sample a batch of single-step transitions $x$ from the replay buffer for policy optimization. \n\nIn principle, one can use the exact same sampling strategy for both context encoder and RL training. Our implementation is meant for easier integration with other baselines. Thank you for your advice, we will make this more clear in the final version. \n\nQ6.  \"it should be impossible to infer task from random-quality dataset\"\n\nFor Ant-Dir, the environments have dense rewards and tasks only differ in reward function. Therefore even with the same state and action distribution, the tasks should still be distinguishable by looking at the reward $r$ conditioning on each $(s, a)$ tuple. Thus the superior performance of UNICORN on random-quality dataset is a strong evidence of its robustness.\n\n[1] Representation learning with contrastive predictive coding\n\n[2] Provably improved context-based offline meta-rl with attention and contrastive learning\n\n[3] Auto-encoding variational bayes", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DtlhKIuA7j", "reviewer": "Reviewer_7Lx7", "summary": "In this paper, the authors propose a unified mathematical framework that encapsulates a subset of the developments conducted in Context-Based Offline Meta-Reinforcement Learning. They achieve this by carefully describing each previous attempt to solve this problem, and how the incremental improvements have been conducted by building upon previously proposed models. Next, they note that all of these previous works can be described as different but related information theory quantities, largely corresponding to the bounds of the (proposed) quantity that should ideally be optimally maximized, corresponding to the mutual information between the latent variables and task variables. They thoroughly explain each component of the theory framework by performing mathematical analysis and giving graph descriptions that enhance comprehension. Based on this framework, the authors propose a method called UNICORN, designed to maximize the ideal mutual information, for supervised and self-supervised settings. Then, they perform experiments under different conditions comparing the proposed method to the previous attempts described in the paper, achieving better results across various experimental settings.", "questions": "1.- In equation 7, what happened to the term $H(M)$? The only explanation I found that makes it disappear is that if you are maximizing mutual information, that term is irrelevant for the maximization since it's fixed, but that doesn’t mean $I(Z;M) = -H(M|Z)$?\n\n2.- The paragraph starting at line 156 discusses CSRO optimizing a convex interpolation of the upper bound $I(Z;X)$ and lower bound $I(Z;X_t, | X_b)$. How is this beneficial compared to just optimizing the lower bound? If the upper bound is\n$$ I(Z;X) = I(Z; X_{t}|X_{b}) + I(Z;X_{b}), $$ then the convex combination of the upper and lower bound would be $$ I(Z; X_{t}|X_{b}) + \\lambda I(Z;X_{b}) $$(also similar to the self supervised objective in eq 12), and according to the framework, the second term includes correlations due to the current policy and not the model of the environment. How optimizing this convex combination would be beneficial? In practice seems like this is a good idea based on the performances of UNICORN and CSRO compared to FOCAL and CORRO.\n\n3.- Are there other methods that try to solve COMRL that are not captured by your framework? What are the differences in their formulation and why can’t they be captured by your framework, and how you would potentially go about integrating these to the theory?\n\n4.- Can this be extended to Online Meta RL? at least from the formalization seems that it could be directly applicable, except for the time-evolving policy which can make $I(Z;X_{b})$ difficult to deal with.", "limitations": "The authors discuss some limitations, and the questions above may address potentially new ones.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper is very well motivated by highlighting previous work contributions to the field and how they fit with their framework.\n- The paper strikes a good balance between exposing technical details and explaining the intuition behind each component of the mathematical analysis. By clearly linking each component to previous work, the authors effectively unify the goals of the models previously mentioned in the paper.\n- By indicating what is the ideal quantity to be maximized, they propose two loss functions, one for supervised and one for self-supervised learning, also explaining how this can be understood in light of previous work.\n- The experiments are well-designed, comparing UNICORN to the relevant and up-to-date models encompassed by its framework. Additionally, UNICORN achieves better results than the baseline while being grounded in a more principled methodology.", "weaknesses": "- The clarity of some aspects of the mathematical derivations and theoretical conclusions could be improved (see specific questions below).\n\n- This is likely a matter of presentation rather than a fundamental issue the approach.\nThe paper would benefit from a clearer explanation of which relevant works from the existing literature are not explicitly incorporated into the proposed framework (see specific questions below).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the constructive comments and the appreciation of our work. Our response to your concerns/questions:\n\nQ1. Eqn 7\n\nEqn 7 concludes that $I(Z; M)\\equiv -H(M|Z)$, which by definition of equivalence $\\equiv$ in Thm 2.3, suggests that the two terms are equal up to a constant, which is indeed $H(M)$. You are correct, $H(M)$ is dropped because it doesn't involve $Z$, and can therefore be regarded as a constant for optimization.\n\nQ2. \"How is convex interpolation beneficial?\"\n\nThe joint distribution of (s, a) depends on several factors: the initial state distribution $\\rho_0$, the transition function $T$ and the behavior policy $\\pi_{\\beta}$. Therefore even it's highly correlated with $\\pi_{\\beta}$, it does in principle contain $\\textit{some}$ causal information about the task identity $M$, especially for tasks that differ in transition dynamics (e.g., Walker, Hopper). Completely dropping the term might be an overkill. We briefly explained this in line 159-161 but will further clarify in the final version. Please check our general reponse for more discussion.\n\nQ3. Scope of the framework\n\nVery interesting question. The level of integration of existing methods to our theory depends on the level of abstraction. Given our Table 1 which summarizes the most represenative COMRL methods to our knowledge, mathematically, we expect our framework to be able to incorporate almost all existing CORML methods that focus on task representation learning. Since according to our Definition 2.1, as long as the method tries to solve CORML by learning a sufficient statistic $Z$ of $X$ w.r.t $M$, it will eventually come down to optimizing an information-theoretic objective equivalent to $I(Z; M)$, or a lower/upper bound like the ones introduced by Thm 2.3, up to some regularizations or constraints. At this level of abstraction, we cannot think of any exception to our best knowledge but are definitely open to discussion.\n\nSince our framework assumes a decoupling of task representation learning and offline policy optimization, UNICORN currently doesn't account for COMRL methods which make nontrivial modification to the standard actor-critic RL learning loop.  This limitation is discussed in line 341-343. For these methods, however, we expect UNICORN to be able to act as a plug-and-play task representation learning module to further improve their task inference robustness and overall performance.\n\nMoreover, if we look at the implementation level, there could be myriads of design choices even for optimizing the same objective. For example, replacing the VAE generator in UNICORN-SS by a diffusion/GAN model might be helpful for image-based tasks, using DT/Diffusion [1] RL backbones might exhibit better scaling for larger datasets and more complex tasks. Also there can be different paradigms, such as the Bayes-Adaptive RL [2][3] which infer $z$ as a belief of the whole history in pursuit of bayes-optimal control, at the cost of breaking the Markovian property of the MDP as well as higher computation. From this perspective, the current supervised and self-supervised UNICORN formulations are far from a complete coverage of all these possible options. Nevertheless, we believe our framework provides a principled guideline for designing novel CORML algorithms from the task representation learning perspective, and can be further extended and improved on a case-by-case basis (e.g., for specific settings or sub-problems of COMRL).\n\nQ4. Extending UNICORN to Online Meta-RL\n\nWe agree that the mathematical formalization should be directly applicable to online RL. However, many of our key derivations rely on the static assumption of $M$ and $X$ (e.g., line 485), which are apparently violated in the online scenario. Extending our framework to online RL is interesting and nontrivial, which might be a good topic for future work.\n\n[1] MetaDiffuser: Diffusion Model as Conditional Planner for Offline Meta-RL\n\n[2] ContraBAR: Contrastive Bayes-Adaptive Deep RL\n\n[3] Varibad: A very good method for bayes-adaptive deep rl via meta-learning", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zbR3q4zl5c", "reviewer": "Reviewer_AbCD", "summary": "This paper provides an analysis of existing Offline Meta-RL (OMRL) methods by considering which informational quantity is maximized by which approach, and how they are related.\nThe show, under a certain assumption on data generation process, that early approaches like FOCAL maximize an upper-bound of the quantity of interest and propose a new method that instead maximizes a lower bound, resulting in a better performance.\nThey also provide an exhaustive evaluation of their approach in multiple different settings (in-domain, ood, model-based, varying data quality).", "questions": "* As mentioned in the weaknesses above, my biggest concern is the interpretation of assumption 2.2, i.e.  $I(Z;M) \\geq I(Z:M|X_b)$. Some additional justification or clarification would be helpful. When does it apply? When is it violated?\n\n* It would be nice to include ContraBAR [Choshen and Tamar; ICML2023] into the comparison of methods, as I believe it also maximizes $I(Z;M)$ in this framework, and is also predictive (by using Contrastive Predictive Coding). It also is applied online and offline. Without it, this submission seems to suggest that it provides the first method maximizing $I(Z;M)$, which I don't believe is correct.\n\n* What insights does Theorem 2.4 provide?\n\n* Are the proposed methods limited to discrete task spaces (during training and execution?)? If so, how could the proposed methods be extended to continuous task spaces?\n\nI'm willing to improve my score if the concerns, especially about Assumption 2.2, are addressed.", "limitations": "* The methods as presented seem to be limited to discrete sets of tasks. Adapting them to continuous state should be feasible, for example by using InfoNCE in UNICORN-Sup, but as such variants are not discussed or evaluated this limitation should be stated.\n\nAt least it should be mentioned in 3.1) when introducing the evaluation tasks.", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "* The information theoretical analysis seems novel, especially in its analysis and comparison of the prior existing methods.\n* The new proposed methods are well motivated based on the prior analysis and perform well on a diverse set of problem setting variations\n* Soure code is provided", "weaknesses": "* The presentation of the Information Theoretical Framework in Section 2.2 is kind of hard to follow. The distinction of \"minimally causally related\" and \"maximally causally related\" quantities seems nonstandard to me and should be clarified. Maybe just writing $I(Z;M) \\geq I(Z:M|X_b)$ is more clear? I would also appreciate a more clear justification of this assumption.\n\n * Experimental validation explores a broad set of settings, but does not discuss hyper-parameter choice and optimization, except for one ablation. A proper hyper-parameter optimization for baselines and proposed methods would significantly strengthen the reliability of the evaluation.\n\n * Theorem 2.4 (Concentration Bound for Supervised UNICORN) is just a straight-forward application of Chebyshev, is not discussed further and doesn't seem to provide any deeper insights. It is thus unclear to me why it was included in the paper, what additional insights it provides.\n\nMinor Points:\n * Figure 1 a) and 1 b) seem unnecessarily complex, a more sketch-like format might be easier to understand\n * It is not clear from the paper how exactly the \"Supervised\" baseline encodes the context.\n * Interaction Information $I(X;Y;Z)$ is not commonly used outside the causal inference community as far as I know, so it would be nice to quickly define it for the reader.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the constructive comments, which really help us make the paper stronger. Our response to your concerns/questions:\n\n# Justification of Assumption 2.2\n\nFor clarification of $I(Z; M; X_b) \\ge 0$ or equivalently $I(Z; M) \\ge I(Z; M|X_b)$, please see our following response to Q1. \"minimally causally related\" and \"maximally causally related\" are indeed qualitative descriptions to convey that the causality between $Z$ and $X_t$ is stronger compared to that between $Z$ and $X_b$. According to the three levels of causation [1], we can define the former as a conditional causality, which means $X_t$ is a $\\textit{necessary}$ factor to infer task identity $Z$, but not sufficient. When conditioning on $X_b$, it provides sufficient information about the reward and transition function, which makes the corresponding $I(Z; X_t|X_b)$ the primary causality.\n\nIn constrast, the causal relation between $Z$ and $X_b$ is much weaker, since the distribution of $X_b$ mostly depends on the behavior policy $\\pi_{\\beta}$, a confounding factor. Therefore the relation should be considered a contributory causality, which means the knowledge about $X_b$ may provide some information about $Z$, but is neither necessary nor sufficient to fully determine $Z$, which induces our definition of $I(Z; X_b)$ as the lesser causality. It's worth to note that quantifying causality in terms of mutual information which we did in our paper is also commonly seen in other literature [2] [3].\n\n# Hyper-Parameters\nWe set the same batch-size, learning rate and network for all context-based methods. The same principle applies for all gradient-based methods and transformer-based method. Additionally, by adjusting the head of the transformer, we ensure that its number of parameters stay close to the policy network in gradient-based method and context-based methods.\n\nFor all methods on each environment, we do thorough grid-search to select the hyper-parameters:\n\nPrompt-DT: $K~[2,5,10,40]$, same as its paper.\n\nMACAW: $\\lambda\\sim [0.005, 0.01, 0.05, 0.1]$, where 0.01 is recommended in its paper.\n\nFOCAL: $\\beta\\sim[1,2,8,16]$ and $n\\sim[-2,-1,1,2]$, same as its paper.\n\nCORRO: we fine-tune the training process of the negative-generator, reducing the training loss from 2.833 (open-source repo) to 0.084.\n\nCSRO: $\\lambda_{FOCAL}\\sim[0.01, 0.03, 0.07, 0.1, 0.3, 0.7, 1.5, 3, 7, 10, 25, 50]$ and $\\lambda_{CLUB}\\sim[0.01, 0.03, 0.07, 0.1, 0.3, 0.7, 1.5, 3, 7, 10, 25, 50]$, where $10,25,50$ are recommended in its paper.\n\nUNICORN-SUP: no hyper-parameter to tune. Hence, UNICORN-SUP is recommended if no hyper-parameter tuning is a priority.\n\nUNICORN-SS: $\\frac{\\alpha}{1-\\alpha}\\sim [0.01, 0.03, 0.07, 0.1, 0.3, 0.7, 1.5, 3, 7, 10, 25, 50]$. Despite our search for UNICORN-SS, we find that it achieves better performance on many parameter choices.\n\nDue to the space limit, we demonstrate the HP records of Prompt-DT, MACAW and CSRO on Ant-Dir in our general response PDF Figure 2.\n\n# Thm 2.4\n\nPlease refer to our following response to Q3.\n\n# How UNICORN-SUP encodes the context\n\nWe use the output feature of the encoder as $z$, which is also the input of the classifier. \n\n# Questions\n\nQ1. \"Interpretation of $I(Z; M)\\ge I(Z; M|X_b)$\"\n\nWe now provide a rigorous proof:\n\n$I(Z; M|X_t, X_b) = 0 \\Longrightarrow I(Z; M) - I(Z; M|X_t, X_b) = I(M; X_b, X_t) - I(M; X_b, X_t|Z) \\ge 0$\n\n$\\Longrightarrow I(M; X_b) + I(M; X_t|X_b) - I(M; X_b|Z) - I(M; X_t|X_b, Z) \\ge 0$\n\n$\\Longrightarrow I(M; X_b) - I(M; X_b|Z) \\ge \\underbrace{-I(M; X_t|X_b)}_{const} + \\underset{\\ge 0}{I(M; X_t|X_b, Z)} \\ge const$\n.\nSince $I(Z; M) - I(Z; M|X_b) =  I(M; X_b) - I(M; X_b|Z)$, we have $I(Z; M) - I(Z; M|X_b) \\ge const$\n\nNamely, for offline setting, $I(Z; M)\\ge I(Z; M|X_b)$ holds up to a constant, which can directly be used to prove the LHS of Thm 2.3 (line 485). We are grateful that your question helps us rethink and rigorize our theory. With the proof above, Assumption 2.2 is no longer an assumption, we will make it a definition of graphical models of COMRL instead in the final version.\n\nQ2. ContraBAR \n\nThank you for the advice. ContraBAR is indeed an interesting paper, which considers contrastive learning for Bayes-Adaptive RL. Please check the pdf in our general reponse for empirical comparison between UNICORN, ContraBAR and BOReL. UNICORN still maintains its edge.\n\nBased on Eqn 3 of the paper, contraBAR optimizes $I(Z; X_t|A)$, which closely resembles the primary causality $I(Z; X_t|X_b) = I(Z; X_t|S, A)$ in our paper and the CORRO loss, but not $I(Z; M)$ directly. Also it follows the paradigm of VariBAD and BOReL which aggregates a long history/trajectory to infer the belief $z$ (therefore non-Markovian), considers online adaptation at test time; in contrast, UNICORN uses single-step transition for task inference (Markovian), and should be applicable to both offline and online adaptation since it's robust to behavior OOD. \n\nQ3. Thm 2.4\n\nPlease see our general response Q4.\n\nQ4.  discrete task spaces\n\nThe 6 MuJoCo and MetaWorld benchmarks used in the current paper are all continuous control environments, so it has not been tested on discrete task spaces. We agree with this limitation. However, extending UNICORN to discrete space should be straightforward: for task representation learning, replace the model backbone to one that can deal with discrete observation, e.g., CNN/ViT for image and transformer/LSTM for language tokens; for downstream RL, replace SAC agent with its variant that can handle discrete actions [4] or other policy networks such as DQN [5]. Such experiments can be addded to the 1 extra page should the paper be accepted.\n\n[1] http://media.acc.qcc.cuny.edu/faculty/volchok/causalMR/CausalMR3.html\n\n[2] Conditional Mutual Information to Approximate Causality for Multivariate Physiological Time Series\n\n[3] On Causality and Mutual Information\n\n[4] Soft actor-critic for discrete action settings\n\n[5] Human-level control through deep reinforcement learning", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank the reviewers, ACs and PCs for ensuring high quality review of the paper. We find all reviews constructive and helpful for making our paper stronger. As requested, the attached PDF provides result of UNICORN vs. a new baseline, ContraBAR, as well as hyper-parameter optimization records of baselines including Prompt-DT, MACAW and CSRO. 

Here we summarize some key/common questions raised and provide our general respose as follows:

Q1. Clarification of Assumption 2.2, in particular $I(Z; M; X_b) \ge 0$ or equivalently $I(Z; M) \ge I(Z; M|X_b)$?

We are grateful to reviewer AbCD who helped us rethink and futher rigorize our theory. This assumption was originally proposed as a key step for proving $I(Z; M) \ge I(Z; X_t|X_b)$ (line 485). However, we now conclude that it can be rigorously proven up to a constant, which is sufficient for deriving Thm 2.3. Please check the detailed proof in our response to reviewer AbCD, Q1.

Q2. About Thm 2.3, why is adding the lesser causality $I(Z; X_b)$ beneficial?

The joint distribution of (s, a) depends on several factors, including the initial state distribution $\rho_0$, the transition function $T$ and the behavior policy $\pi_{\beta}$. Therefore it contains two components: the correlation with $\pi_{\beta}$, which is spurious; the correlation with $\rho_0$ and $T$, which are contributory causalities [1] that capture partial information about $M$. Therefore, optimizing the primary causality $I(Z; X_t|X_b)$ alone eliminates the spurious correlation at the cost of overlooking the contributory causalities, which is likely an overkill. Instead, an ideal choice would be a trade-off between the spurious correlation and the contributory causalities, effectively what UNICORN-SS and CSRO do, resulting in better empirical performance.

Q3. Scope of UNICORN: Are there any COMRL methods not captured by the framework? If yes, why and how to incorporate them?

The level of integration of existing methods to our theory depends on the level of abstraction. Given Table 1 which summarizes the most represenative COMRL methods to our knowledge, mathematically, we expect our framework to be able to incorporate almost all existing COMRL methods that focus on task representation learning. Since according to our Definition 2.1, as long as the method tries to solve COMRL by learning a sufficient statistic $Z$ of $X$ w.r.t $M$, it will eventually come down to optimizing an information-theoretic objective equivalent to $I(Z; M)$, or its lower/upper bound like the ones introduced by Thm 2.3, up to some regularizations or constraints. At this level, we are unaware of any exception but are definitely open to discussion.

Since our framework assumes a decoupling of task representation learning and offline policy optimization, UNICORN currently doesn't account for COMRL methods which make nontrivial modification to the standard actor-critic RL learning loop. This limitation is discussed in line 341-343. For these methods, however, we expect UNICORN to be able to act as a plug-and-play task representation learning module to further improve their task inference robustness and overall performance.

Moreover, if we look at the implementation level, there could be myriads of design choices even for optimizing the same objective. For example, replacing the VAE generator in UNICORN-SS by a diffusion/GAN model might be helpful for image-based tasks, using DT/Diffusion [2] backbones might exhibit better scaling for larger datasets and more complex tasks. Also there can be different paradigms, such as the Bayes-Adaptive RL [3][4]. From this perspective, the current supervised and self-supervised UNICORN formulations are far from a complete coverage all these possible choices. Nevertheless, we believe our framework provides a principled guideline for designing novel COMRL algorithms from the task representation learning perspective, and can be further extended and improved on a case-by-case basis (e.g., specific settings or sub-problems of COMRL).

Q4. Significance of Thm 2.4.

The interpretation of Thm 2.4 is only briefly discussed in line 344-345 due to the page limit. The main purpose of Thm 2.4 is to show that, even though UNICORN-SUP seems to directly optimize $I(Z;M)$ on paper, it's not perfect by making a finite-sample approximation using limited number of task instances. For our experiments, $n_M\le 40$, which may induce a significant approximation error according to Eqn 11. This explains why in most of our experiments, UNICORN-SUP is inferior to UNICORN-SS. We believe that with more tasks and data, its performance can be improved. 

Q5. Given the mathematical resemblence/equivalence of COMRL methods, why does UNICORN achieves suprior performance?

Compared to FOCAL and CORRO, directly optimizing $I(Z; M)$ or a convex combination of the primary and lesser causalities enforce a better trade-off between spurious correlation and contributory causality, see our response to Q2.
 
The comparison between UNICORN and CSRO is trikier, since the objective of CSRO and UNICORN-SS are mathematically equivalent. We believe the main difference lies in the implementation choices. CSRO approximates $I(Z; X_b)$ with the CLUB loss. On top of this approximation, it employs variational distribution $q_{\theta}(z|x_b)$ to approximate the true posterior $q(z|x_b)$ in Eqn 6. Therefore it involves $\textit{2 levels}$ of approximation. In contrast, both UNICORN variants involves only ${1 level}$ of approximation (finite sample approximation in Thm 2.4 for UNICORN-SUP and variational approximation in eqn 14 for UNICORN-SS). We believe this explains (at least partially) the superiority of UNICORN over CSRO.

Q6. Additional information (e.g., proof sketch, interpretation of Thm 2.4, experiments on discrete spaces, more related work...)

We appreciate all the suggestions and will provide these information using the 1 extra page should the paper be accepted.

Meta-review: This paper provides a theoretical analysis that explains various offline meta-RL approaches in a unified way and proposes a new algorithm that outperforms the baselines across a broad range of offline meta-RL benchmarks. All of the reviewers unanimously found the theoretical analysis novel and insightful, as it unifies seemingly different approaches. They also found that the empirical results are strong and convincing enough. I believe that the findings of this paper would be interesting to the meta-RL community. Thus, I recommend to accept this paper.
Decision: Accept (spotlight)