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_cc_2024_Conference
Published: 2024-05-11 16:39:43

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.

Main paper: 
 introduction: 
 The ability to swiftly learn and generalize to new tasks is a hallmark of human intelligence. In pursuit of this high level of artificial intelligence (AI), the paradigm of meta-reinforcement learning (RL) proposes to train AI agents in a trial-and-error manner by interacting with multiple external environments. In order to quickly adapt to the unknown, the agents need to integrate prior knowledge with minimal experience (namely the context) collected from the new tasks or environments, without over-fitting to the new data. This meta-RL mechanism has been adopted in many applications such as games [1, 2] , robotics [3, 4] 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-critical scenarios such as healthcare [8] and robotic manipulation [9] . As a remedy, offline RL [10] enables agents to learn from logged experience only, therefore circumventing risky or costly online interactions. Recently, offline meta-RL (OMRL) [11] [12] [13] has emerged as a novel paradigm to significantly extend the applicable range of RL by "killing two birds in one stone": it builds powerful agents that can quickly learn and adapt by meta-learning, while leveraging offline RL mechanism to ensure a safe and efficient optimization procedure. In the context of classical supervised or self-supervised learning, which is de facto offline, OMRL is reminiscent of the multi-task learning [14] , meta-training [6] and fine-tuning [15] [16] [17] of pre-trained large models. We envision it as a cornerstone of RL foundation models [18] [19] [20] in the future.Along the line of OMRL research, context-based offline meta-reinforcement learning (COMRL) is a popular paradigm that seeks optimal meta-policy conditioning on the context of Markov Decision Processes (MDPs). Intuitively, the crux of COMRL 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 [11] proposes to capture the structure of task representations by distance metric learning. From a geometric perspective, it 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.Despite its effectiveness, FOCAL is reported to be vulnerable to context shifts [21] , i.e., when testing on out-of-distribution (OOD) data (Fig. 1 ). 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 [21, 22] . To alleviate the problem, follow-up works such as CORRO [23] reformulates the task representation learning of COMRL as maximizing the mutual information I(Z; M ) between the task variable M and its latent Z. It then approximates I(Z; M ) by an InfoNCE [24] contrastive loss, where the positive and negative pairs are conditioned on the state-action tuples (s, a). Inspired by CORRO, a recently proposed method CSRO [22] 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 [25] benchmarks.Contributions In this paper, following the recent development and storyline of COMRL, we present a Unified Information Theoretic Framework of Context-Based Offline Meta-Reinforcement 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 I(Z; M ) and their linear interpolation respectively, which provides a 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 1 ) concerning their objectives and implementations, we highlight the potential design choices of novel algorithms offered by our framework. As examples, we investigate two instantiated algorithms, one supervised and the other self-supervised, and demonstrate 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 approximations/regularizations of I(Z; M ), as well as new implementations to further combat context shift. 
preliminaries, problem statement and related work: 
 We consider MDP modeled as M = (S, A, T, R, ρ 0 , γ, H) with state space S, action space A, transition function T (s ′ |s, a), bounded reward function R(s, a), initial state distribution ρ 0 (s), discount factor γ ∈ (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 π, then the environment updates the state with T and returns a reward with R. We denote the marginal state distribution at time t as µ t π (s). The V-function and Q-function are given by 2), which is also a challenging scenario for COMRL models to generalize.EQUATIONEQUATIONIn this paper, we restrict to scenarios where tasks share the same 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 i = (S, A, T i , R i , ρ 0 , γ, H) ∈ M, where M is the set of all possible MDPs to be considered. For each task index i, the offline dataset X i = {(s i j , a i j , r i j , s ′ i j )} is collected by a behavior policy π i β . For meta-learning, given a trajectory segment c i 1:n = {(s i j , a i j , r i j , s ′ i j )} n j=1 of length n as the context of M i , COMRL employs a context encoder q ϕ (z|c i 1:n ) to obtain the latent representation z i of task M i . If z contains sufficient information about the task identity, COMRL can be treated as a special case of RL on partially-observed MDP [26] , where z is interpreted as a faithful representation of the unobserved state. By conditioning on z, the learning of universal policy π θ (•|s, z) and value function V π (s, z) [27] become regular RL, and optimality can be attained by Bellman updates with theoretical guarantees. To this end, FOCAL [11] 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:EQUATIONHowever, it is empirically shown that FOCAL struggles to generalize in the presence of context shift [21] . This problem has been formulated in several earliest studies of OMRL. Li et al. [28] observed that when there are large divergence of the state-action distributions among the offline datasets, due to shortcut learning [29] , the task encoder may learn to ignore the causal information like rewards and spuriously correlate primarily state-action pairs to the task identity, leading to poor generalization performance. Dorfman et al. [13] concurrently identified a related problem which they termed MDP ambiguity in the context of Bayesian offline RL. An example is illustrated in Figure 1 . The problem is further exacerbated in the fully offline setting, as the testing distribution is fixed and cannot be augmented by online exploration, allowing no theoretical guarantee for the context shift.To mitigate context shift, a subsequent algorithm CORRO [23] proposes to optimize a lower bound of I(Z; M ) in the form of an InfoNCE [24] contrastive lossEQUATIONwhere x = (s, a, r, s ′ ), z ∼ q ϕ (z|x), h(x, z) = P (z|x) P (z) ≈ q ϕ (z|x) P (z) , and x * = (s, a, r * , s ′ * ) as a transition tuple generated for task M * ∈ M conditioned on the same (s, a). To further combat the spurious correlation between the task representation and behavior-policy-induced state-actions, a recently proposed method CSRO [22] introduces a CLUB upper bound [30] of the mutual information between z and (s, a), to regularize the FOCAL objective:EQUATIONEQUATIONIn the next section, we will show how these algorithms are inherently connected and how their context-shift robustness gets improved incrementally.We start with a formal definition of task representation learning in COMRL: Definition 2.1. Given an input context variable X ∈ X and its associated task/MDP random variable M ∈ M, task representation learning in COMRL aims to find a sufficient statistics Z of X with respect M .In pure statistical terms, Definition 2.1 implies that an ideal representation Z is a mapping of X that captures the mutual information I(X; M ). We therefore define the following dependency structures in terms of directed graphical models: Definition 2.2. The dependency graphs of COMRL is given by Fig. 2 , where X b and X t are the behavior-related (s, a)-component and task-related (s ′ , r)-component of the context X, with X = (X t , X b ). For the first graph, M → X → Z forms a Markov chain, which satisfies I(Z; M |X) = 0. To interpret the second graph, given an MDP M ∼ M , the state-action component of X is primarily captured by the behavior policy π β : s ∼ µ π β (s), a ∼ π β . The only exception is when tasks differ in initial state distribution ρ 0 or transition dynamics T , in which case the state variable S also depends on M . We therefore define it as the behaviorrelated component, which should be minimally causally related (dashed lines) to M and Z. Moreover, when X b is given, X t is fully characterized by the transition function T : (s, a) → s ′ and reward function R : (s, a) → r of M , which should be maximally causally related (solid lines) to M and Z and therefore be defined as the task-related component. Accordingly, we name I(Z; X t |X b ) and I(Z; X b ) the primary and lesser causality in our problem respectively. With this prior knowledge and assumption, we present the central theorem of this paper: Theorem 2.3. Let ≡ denote equality up to a constant, then Secondly, CSRO as the latest COMRL algorithm, inherently optimizes a linear combination of the FOCAL and CORRO objectives. In the 0 ≤ λ ≤ 1 regime, the CSRO objective becomes a convex interpolation of the upper bound I(Z; X) and the lower bound I(Z; X t |X b ) of I(Z; M ), which in essence, enforces a trade-off between the causal (Z with T & ρ 0 ) and spurious (Z with π β ) correlation contained in I(Z; X b ). This accounts for the improved performance of CSRO compared to FOCAL and CORRO.I(Z; X t |X b ) primary causality ≤ I(Z; M ) ≤ I(Z; X t |X b ) + I(Z; X b ) = I(Z; X)By providing a unified view of pre-existing COMRL algorithms, Theorem 2.3 opens up 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; M ) can be re-expressed asEQUATIONwhere H(•) is entropy. Since in practice, each z i of sample x i is collected within a specific task M i , minimizing the parameterized entropy H θ (M |Z = z i ) is equivalent to finding an optimal function p θ (M |z) which correctly assigns the ground-truth label M i to z i , i.e., optimizingEQUATIONSuppose the task label M is given during training and a total of n M training tasks {M i } n M i=1 are drawn from the task distribution p(M ) for meta-training. Under this supervised scenario, by substituting Eqn 8 into 7, we haveEQUATIONwhich is precisely the negative cross-entropy loss H(M , P (M |X)) for n M -way classification with feature z and classifier p θ . We therefore define the objective of supervised UNICORN asEQUATION)Note that L UNICORN-SUP is convex and operates as a finite-sample approximation of -I(Z; M ), for which we give the following theorem: Theorem 2.4 (Concentration bound for supervised UNICORN). Denote by Î(Z; M ) the empirical estimate of I(Z; M ) by n M tasks, Ī(Z; M ) the expectation, then with probability at least 1 -δ,EQUATION)Proof. See Appendix B. Self-supervised UNICORN In practice, offline RL datasets may often be collected with limited knowledge of the task specifications or labels. In this scenario, previous works implement self-supervised learning [31] to obtain effective representation Z, such as the contrastive-based FOCAL/CORRO to optimize I(Z; X)/I(Z; X t |X b ) respectively; or generative approaches like VariBAD [32] /BOReL [13] to reconstruct the trajectories X by variational inference, which is equivalent to maximizing I(X t ; Z, X b ). By Theorem 2.3, these methods optimize a relatively loose upper/lower bound of I(Z; M ), which can be improved by a convex combination of the two bounds:EQUATION)where 0 ≤ α ≤ 1 is a hyperparameter. Implementing each term in Eqn 12 allows ample decision choices, such as the contrastive losses in Eqs. ( 3), ( 4) and ( 6) or autoregressive generation via Decision Transformer [33, 34] or RNN [32] . For demonstration, in this paper we employ a contrastive objective L FOCAL as in Eq. ( 3) for I(Z; X) while approximate I(Z; X t |X b ) by reconstruction. By the chain rule of mutual information [35] :EQUATIONsince I(X t ; X b ) is a constant when X t and X b are drawn from a fixed distribution as in offline RL. Moreover, by definition of mutual information:EQUATIONwhich induces a generative objective L recon := -I(X t ; Z, X b ) by reconstructing X t with a decoder network p θ (•|z, x b ) conditioning on Z and X b . As a result, the proposed unsupervised UNICORN objective can be rescaled as Eqn 15. The influence of the hyper-parameter α 1-α is shown in Appendix C. 4 .EQUATION)We summarize our learning procedure in Fig. 3 with pseudo-code in Algorithms 1 and 2. A holistic comparison of our proposed algorithms with related contextual meta-RL methods is shown in Table 1 .The extra KL divergence can be interpreted as the result of a variational approximation to an information bottleneck [36, 37] that constrains the mutual information between Z and X, which we found unnecessary in our offline setting (see ablation in Table 6 .) Behavior regularized actor critic [38] is employed to tackle the bootstrapping error [39] for downstream offline RL implementation. 
experiments: 
 For brevity, we name our proposed supervised and self-supervised algorithms UNICORN-SUP and UNICORN-SS respectively. To evaluate their effectiveness, our main experiments are organized to [11, 21] Offline I(Z; X) Contrastive Transition CORRO [23] Offline [40] Offline I(X t ; Z, X b ) Generative Transition BOReL [13] OfflineI(Z; X t |X b ) Contrastive Transition CSRO [22] Offline (1 -λ)I(Z; X) + λI(Z; X t |X b ) Contrastive Transition GENTLEI(X t ; Z, X b ) -D KL (q ϕ (Z|X)||p θ (Z)) Generative Trajectory VariBAD [32] Online I(X t ; Z, X b ) -D KL (q ϕ (Z|X)||p θ (Z)) Generative Trajectory PEARL [7] Online -D KL (q ϕ (Z|X)||p θ (Z)) N/A Transition ContraBAR [41] Offline&Online I(Z; X t |A) Contrastive Trajectoryaddress three primary inquiries regarding the core advantages of UNICORN: (1) 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.We adopt MuJoCo [25] and MetaWorld [3] benchmark to evaluate our method, involving six robotic locomotion tasks which require adaptation across reward functions (HalfCheetah-Dir, HalfCheetah-Vel, Ant-Dir, Reach) or across dynamics (Hopper-Param, Walker-Param). We compare UNICORN with six competitive OMRL algorithms. These baselines are categorized as context-based, gradientbased and transformer-based methods:FOCAL, CORRO, CSRO are context-based methods that can be seen as special cases or approximations of UNICORN, see details in Section 2.Supervised is a context-based method, directly using actor-critic loss to train the policy/value networks and task encoder end-to-end. We find it a competitive baseline across all benchmarks.MACAW [12] is a gradient-based method, using supervised advantage-weighted regression for both the inner and outer loop of meta-training.Prompt-DT [42] is a transformer-based method, taking context as the prompt of a Decision Transformer (DT) [33, 34] to solve OMRL as a conditional sequence modeling problem.We perform our evaluation in a few-shot manner. For gradient-based methods, we first update the meta-policy with a batch of testing context and then evaluate the agent. For transformer-based methods, we utilize the trajectory segment as the prompt to condition the auto-regressive rollout. -SS 1307±26 1296±24 -22±1 -94±5 267±14 236±18 316±6 304±11 419±44 407±46 2775±241 2604±183 UNICORN-SUP 1296±20 1130±76 -25±3 -91±5 250±4 239±16 312±4 302±12 322±28 312±39 2681±111 2641±140 CSRO 1180±228 458±253 -28±1 -102±5 276±19 233±12 310±6 301±10 310±58 279±65 2720±235 2801±182 CORRO 704±450 245±146 -37±3 -112±2 148±13 120±12 283±8 272±13 277±38 213±48 2468±175 2322±327 FOCAL 1186±272 861±253 -22±1 -97±2 217±29 173±24 302±4 297±13 308±98 286±91 2424±256 2316±303 Supervised 962±356 782±429 -24±1 -104±1 238±39 202±38 306±10 294±8 256±60 210±28 2489±248 2283±205 MACAW 1155±10 450±6 -56±2 -188±1 26±3 0±0 218±6 205±2 141±9 130±5 2431±157 1728±79 Prompt-DT 1176±40 -25±9 -118±66 -249±21 1±0 0±0 234±5 202±5 185±9 156±17 2165±85 1896±111Denote by {π i,t β } i=1:N M ,t=0:N the checkpoints of behavior policies for all tasks logged during data collection (i.e., training SAC agents), where i refers to the task index and t represents training iteration. In the in-distribution (IID) setting, for each test task M j , we use behavior policies only from the same tasks {π j,t β } j,t=0:N to collect the context and the conditioned rollouts, as visualized in Figure 1 . Figure 4 illustrates the learning curves of UNICORN vs. all baselines, 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 is comparable to UNICORN is expected since it also optimizes a linear combination of the lower and upper bound of I(Z; M ) by Theorem 2. To evaluate OOD generalization, for each test task M j , we use all behavior policies in {π i,t β } i=1:N M ,t=0:N to collect the OOD contexts and the conditioned rollouts, as visualized in Figure 1 . The OOD testing performance is measured by averaging returns across all testing tasks, as shown in the OOD entries of Table 2 . While all methods suffer notable decline when facing OOD context, UNICORN maintains the strongest performance by a large margin. Another observation is that the context-based methods are in general significantly more robust compared to the gradientbased MACAW and transformer-based Prompt-DT. This also delineates the storyline of COMRL development from 2021 to dateFOCAL (2021) → CORRO (2022) → CSRO (2023) → UNICORN (2024)as a roadmap for pursuing more robust and generalizable task representation learning for COMRL.To test whether the UNICORN framework can be applied to different data distributions, we collect three types of datasets whose size is equal to the mixed dataset used in Section 3.2 and 3.3: random, medium and expert, which are characterized by the quality of the behavior policies.As shown in Table 3 , UNICORN demonstrates SOTA performance on all types of datasets. Notably, despite CSRO having better IID results than UNICORN on Ant-Dir under the mixed distribution, UNICORN surpasses CSRO by a large margin under the narrow distribution (medium and expert) and random distribution. -SS 81±18 62±6 220±23 243±10 279±10 262±13 UNICORN-SUP 75±15 60±5 140±11 126±32 247±15 229±19 CSRO 2±3 0±1 166±10 198±17 252±39 202±45 CORRO 1±1 0±0 8±5 -7±2 -4±10 -14±9 FOCAL 67±26 44±10 171±84 187±86 229±42 246±20 Supervised 65±6 47±12 149±50 110±80 249±33 215±60 MACAW 3±1 0±0 28±2 1±1 88±43 1±1 Prompt-DT 1±0 0±0 2±4 0±1 78±15 1±2We observe that CORRO fails on all three datasets. This may be due to the reliance of CORRO on the negative sample generator. Under narrow data distribution or poor random data distribution, it might be more challenging to generate negative samples of good quality. As for Prompt-DT, its performance improves significantly with increased data quality, which is expected since the decision transformer adopts a behavior-cloning-like supervised learning style. 
discussion: 
 This section presents more empirical evidence on the applicability of the UNICORN framework. As UNICORN tackles task representation learning from a general information theoretic perspective, it is in principle model-agnostic. Hence a straightforward idea is to transfer UNICORN to other model architectures like DT, which we envision to be a promising backbone for large-scale training of RL foundation models.We employ a simple implementation by embedding the task representation vector z as the first token in sequence to prompt the learning and generation of DT. Instead of the unlearnable prompt used in the original Prompt-DT, we enforce task representation learning of the prompt by optimizing L UNICORN-SS , L UNICORN-SUP and L FOCAL . We name these variants as UNICORN-SS-DT, UNICORN-SUP-DT and FOCAL-DT respectively.As shown in Table 4 and Figure 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 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 [43] and DT [44] . We leave this verification to future work.With decoder p θ (x t |z, x b ), UNICORN-SS can potentially enable a model-based paradigm by generating data using customized latent z, which can be especially useful for generalization to OOD tasks. We implement this idea on Ant-Dir by sorting the tasks according to the polar angle of their goal direction (θ ∈ [0, 2π)), and construct task-level OOD by taking 20 consecutive tasks in the middle for training and the rest for testing (Figure 1 ). Gaussian noise ϵ is applied to the real task representation z and (s, a, z +ϵ) is then fed to the decoder to generate (s ′ , r). The imaginary rollouts {(s t , a t , s ′ t , r t )} n t=1 are used during the training of RL agent. We adopt the conventional ensemble technique in model-based RL [45, 46] to stabilize the training process.As shown in Figure 5 , UNICORN-SS-enabled model-based RL is the only method to achieve positive performance in this extremely challenging scenario where context shift is induced by OOD tasks.  
conclusion & limitation: 
 In this paper, we unify three major context-based offline meta-RL algorithms, FOCAL, CORRO and CSRO (potentially a lot more) into a single framework from an information theoretic representation learning perspective. Our theory offers valuable insight on how these methods are inherently connected and incrementally evolved in pursuit of better generalization against context shift. Based on the proposed framework, we instantiate two novel algorithms called UNICORN-SUP 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-online COMRL paradigms. Mathematically, we expect our framework to be able to incorporate almost all existing COMRL methods that focus on task representation learning. Since according to Definition 2.1, as long as the method tries to solve COMRL by learning a sufficient statistic of w.r.t , 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 Theorem 2.3, up to some regularizations or constraints.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 at most 40 tasks and ∼ 450k transitions, which might be the reason why UNICORN-SUP is inferior to UNICORN-SS (n M too small in Theorem 2.4). However, we expect our conclusion to 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), which are apparently violated in the online scenario. Extending our framework to online RL is interesting and nontrivial, which we leave for future work. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 to a 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, or as supplemental material. 
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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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 to particular applications, let alone deployments. However, if there is a 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 is 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: 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, or scraped datasets)?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 a 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 a 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? 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 or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.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 package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. 
 
 Appendix: Require: Offline Datasets X = {X i } Ntrain i=1 , training tasks M = {M i } Ntrain i=1 , initialized learned policy π ω , Q function Q ψ , context encoder q ϕ , classifier/decoder p θ , hyper-parameter α, task batch size N tb , learning rates β 1 , β 2 , β 3 , β 4 1: while not done do 2:Sample task batch {M j } N tb j=1 ∼ M and the corresponding replay buffer R = {X j } N tb j=1 ∼ X 3:for step in each iter do 4:// Train the context encoder and decoder 5:Sample context {c j } N tb j=1 ∼ R 6:Obtain task embeddings {z j ∼ q ϕ (z|c j )} N tb j=1 7:Estimate L UNICORN (L UNICORN-SUP or L UNICORN-SS )8:ϕ ← ϕ -β 1 ∇ ϕ L UNICORN 9: θ ← θ -β 2 ∇ θ L UNICORN 10:// Train the actor and critic 11:Detach the task embeddings {z j } N tb j=1 12:Sample training data {d j } N tb j=1 ∼ R 13:Estimate L actor , L critic 14:ω ← ω -β 3 ∇ ω L actor 15: ψ ← ψ -β 4 ∇ ψ L critic 16:end for 17: end while Algorithm 2 UNICORN Meta-testing Require: Offline Datasets X = {X i } Ntest i=1 , testing tasks M = {M i } Ntest i=1 , trained policy π ω and context encoder q ϕ 1: for each M i do 2:Sample context c i ∼ X iObtain task embedding z i ∼ q ϕ (z|c i ) I(Z; M |X t , X b ) = 0 =⇒I(Z; M ) -I(Z; M |X t , X b ) = I(M ; X b , X t ) -I(M ; X b , X t |Z) ≥ 0 =⇒I(M ; X b ) + I(M ; X t |X b ) -I(M ; X b |Z) -I(M ; X t |X b , Z) ≥ 0 =⇒I(M ; X b ) -I(M ; X b |Z) ≥ -I(M ; X t |X b ) const + I(M ; X t |X b , Z) ≥0 ≥ const. Now, since 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 ) = I(Z; M ; X b ) ≥ const ≡ 0.With Lemma B.1, we proceed to prove Theorem 2.3 as follows:Proof. I(Z; X t |X b ) ≡ -L CORRO ≤ I(Z; M ) is shown in the Theorem 4.1. of CORRO [23] and GENTLE [40] with strong assumptions. We hereby present a proof with no assumptions. Given I(Z; M ; X b ) ≥ 0 (Lemma B.1), we haveI(Z; M ) = I(Z; M |X b ) + I(Z; M ; X b ) ≥ I(Z; M |X b ) = I(M ; Z, X t |X b ) -I(M ; X t |Z, X b ) = I(Z; M |X t , X b ) 0 by M →X→Z +I(M ; X t |X b ) -I(M ; X t |Z, X b ) = I(M ; X t |X b ) -I(M ; X t |Z, X b ) = I(M ; X t |X b ) -H(X t |Z, X b ) + H(X t |M , Z, X b ) ≥ I(M ; X t |X b ) -H(X t |Z, X b ) = I(M ; X t |X b ) const -H(X t ) const +I(X t ; Z, X b ) ≡ I(X t ; Z, X b ) = I(Z; X t |X b ) + I(X t ; X b ) const ≡ I(Z; X t |X b ),as desired. The third and fourth lines are obtained by direct application of the chain rule of mutual information [35] . All mutual information terms without Z are held constant in the fully offline scenario.I(Z; M ) ≤ I(Z; X) is a direct consequence of the Data Processing Inequality [35] and the Markov chain M → X → Z in Definition 2.2. We now prove the three claims regarding FOCAL, CORRO and CSRO:1. L FOCAL ≡ -I(Z; X).I(Z; X) := E x,z log p(z, x) p(z)p(x) = E x,z   log   1 p(z) p(z|x) |M|     + log(|M|) = E x,z   log   1 p(z) p(z|x) |M|E x p(z|x) p(z)     + log(|M|) ≈ E x,z   log   1 p(z) p(z|x) M i ∈M E x i ∼X i p(z|x i ) p(z)     + log(|M|) = E x,z   log   p(z|x) p(z) M * ∈M E x i ∼X i p(z|x i ) p(z)     + log(|M|) = E x,z log h(x, z) M i ∈M h(x i , z) + log(|M|)The first term on RHS is precisely the supervised contrastive learning objective introduced by a variant of FOCAL [21] , which is equivalent to the negative distance metric learning loss L FOCAL with the effect of pushing away embeddings of different tasks while pulling together those from the same task. Therefore we haveI(Z; X) = E x,z log h(x, z) M i ∈M h(x i , z) + const ≡ -L FOCAL . 2. L CORRO ≡ -I(Z; X t |X b ). I(Z; X t |X b ) := E xt,x b ,z log p(z, x t |x b ) p(z|x b )p(x t |x b ) = E xt,x b ,z   log   1 p(z|x b ) p(z|xt,x b ) |M|     + log(|M|) = E xt,x b ,z   log   1 p(z|x b ) p(z|X) |M|E M * ∈M p(z|x * t ,x b ) p(z|x b )     + log(|M|) ≈ E xt,x b ,z   log   1 p(z|x b ) p(z|xt,x b ) M * ∈M p(z|x * t ,x b ) p(z|x b )     + log(|M|) = E xt,x b ,z   log   p(z|xt,x b ) p(z|x b ) M * ∈M p(z|x * t ,x b ) p(z|x b )     + log(|M|) = E xt,x b ,z   log   p(z|xt,x b ) p(z) M * ∈M p(z|x * t ,x b ) p(z)     + log(|M|) = E x,z log h(x, z) M * ∈M h(x * , z) + log(|M|) ≡ -L CORRO . 3. L CSRO ≥ (λ -1)I(Z; X) -λI(Z; X t |X b ).The CLUB loss in Eqn 6 operates as a upper bound of I(Z; X b ) [30] . By Eqn 5 we haveL CSRO = L FOCAL + λL CLUB ≥ L FOCAL + λI(Z; X b ) = L FOCAL + λ (I(Z; X) -I(Z; X t |X b )) chain rule of mutual information 1 = -I(Z; X) + λ(I(Z; X) -I(Z; X t |X b )) = (λ -1)I(Z; X) -λI(Z; X t |X b ).Proof. Denote by n Z := n M i=1 |M i ||X i | the number of context samples, we haveEQUATION≃ -EQUATION≃ -EQUATIONSince n Z ≫ n M , the concentration characteristic of I(Z; 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,EQUATION)Or equivalently, with probability at least 1 -δ,EQUATIONRemark: Consider a practical implementation of p(z|M ) as a multivariate Gaussian N (µ M , Σ M ), which givesH(Z|M ) = 1 2 log[(2πe) dim Z |Σ M |] [35]. Substituting into Eqn 20, we have with probability at least 1 -δ,EQUATIONC Further Experiments C.1 UNICORN-SS-0: A Label-free Version When α = 0, UNICORN-SS task representation learning reduces to the minimization of the reconstruction loss only, which becomes a label-free algorithm (i.e., it does not require the knowledge of task identities/labels to optimize, as opposed to the contrastive learning in previous methods). We name this special case UNICORN-SS-0, which is equivalent to a concurrent method GENTLE [40] .For a fair comparison, we use a variant of a Bayesian OMRL method BOReL [13] that does not utilize oracle reward functions, as a label-free baseline. At test time, we modify BOReL to use offline datasets rather than the online data collected by interacting with the environment to infer the task information 2 . As shown in Table 5 , UNICORN-SS-0 is also competitive with BOReL.Table 5 : UNICORN-SS-0 compared to another label-free COMRL method, BOReL, on Ant-Dir.Algorithm Ant-Dir IID OOD UNICORN-SS-0 220±16 200±9 BOReL 206±18 187±10To further validate that the KL constraint D KL (q θ (z|x)|N (0, I)) is unnecessary, we weight the KL constraint to UNICORN-SS-0. Table 6 shows that KL constraint would reduce the performance under the offline setting.For better interpretation, we visualize the task representations by 2D projection of the embedding vectors via t-SNE [47] . For each test task, we generate a trajectory using each policy in {π i,t β } i,t and infer their task representations z. This setting aligns with our OOD experiments in the main text.Since the UNICORN-SS objective operates as a convex combination of the reconstruction loss (UNICORN-SS-0) and the FOCAL loss, we compare representations of UNICORN-SS to these two extremes. As shown in Figure 6 , compared to UNICORN-SS-0, 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. To illustrate the effect of hyper-parameter α 1-α on the asymptotic performance, we set the following ablation study. Figure 8 shows that as α 1-α increases, the performance gradually increases but when it goes excessively the performance would decrease (FOCAL means α 1-α → ∞), which validates our proposed theory. We also find that UNICORN-SS can maintain relatively stable performance over a range of α 1-α values, thus alleviating the exhaustion from parameter-tuning.Table 7 lists the necessary hyper-parameters that we used to produce the experimental results. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.Experiment shows that the usage of offline dataset rather than online data has little effect on the results of BOReL. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: My weaknesses are relatively minor.

# Clarity
Legends 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.

# Significance
I 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.
questions: 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?
2. 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
3. In Table 1, column Context  X. What does "Transition" vs "Trajectory" mean?
4. 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?
5. 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.
6. 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: Limitations are stated and clear.
paper_summary: paper_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.
strength_weakness: strength_weakness: - The clarity of some aspects of the mathematical derivations and theoretical conclusions could be improved (see specific questions below).

- This is likely a matter of presentation rather than a fundamental issue the approach.
The 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).
questions: 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)$?

2.- 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
$$ 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.

3.- 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?

4.- 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: limitations: The authors discuss some limitations, and the questions above may address potentially new ones.
paper_summary: paper_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.
The 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.
They also provide an exhaustive evaluation of their approach in multiple different settings (in-domain, ood, model-based, varying data quality).
strength_weakness: strength_weakness: * 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.

 * 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.

 * 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.

Minor Points:
 * Figure 1 a) and 1 b) seem unnecessarily complex, a more sketch-like format might be easier to understand
 * It is not clear from the paper how exactly the "Supervised" baseline encodes the context.
 * 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.
questions: 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?

* 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.

* What insights does Theorem 2.4 provide?

* 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?

I'm willing to improve my score if the concerns, especially about Assumption 2.2, are addressed.
limitations: 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.

At least it should be mentioned in 3.1) when introducing the evaluation tasks.

Comments: title: Thank You
Comment by Reviewer sHdF
Further Clarification
Thank You
Rebuttal
References
comment: Thank you very much for the constructive comments. We highly appreciate the fruitful discussion, which will undoubtedly improve the quality of our paper. We will make sure to properly implement all reviewers' suggestions in the final version.
I thank the authors for their detailed response. Indeed my concerns are generally addressed. I think the agreed points to be clarified will make a difference. I also hope some of the points which I was corrected on above will be used to improve the clarity and avoid similar confusions in a revised draft. I thank the authors for the nice paper and will raise my score to 8.
Thank you for the comments. The "context training size" refers to the number of transitions used for each task during $\textit{one forward pass}$ (basically we randomly sampe 1 trajectory for task inference). The total number of transitions used to train a single task is given by the "dataset size in Table 7", and is much larger. For example, for Ant-Dir, it's $1e5=100,000$. We will clarify this in the final version.

We agree that if we increase the number of tasks without expanding the dataset proportionally, it might be problematic not just for task representation learning, but also for downstream offline RL since it may suffer from insufficient dataset coverage [1] for each task. However, the hope is that, just like pretraining large language models, we are able to collect decent amount of data for each individual task, and use offline meta-RL as a multi-task pretraining technique to produce robust task representations for a generalist agent to solve decision making problems. This is why we are interested in DT since transformer is a promising backbone with better scaling in the large data regime. 

Again thank you very much for the fruitful discussion. Please do not hesitate to let us know if you have further comments/questions.

[1] What are the Statistical Limits of Offline RL with Linear Function Approximation?
Thank you for your comments and suggestions, which will undoubtedly improve the quality of our paper. We will make sure to properly address your questions in the final version.
Thank you for the reply. I appeciate the clarification, but I made a mistake in my prior comment.

The limitation I was interested in is **during training**. 

Supervised Unicorn is using an N-way classification loss. This probably works very well with discrete sets of tasks, i.e. the N=20 training tasks sampled from the semicircle in Figure 1. If I'm interpreting the Appendix correctly, you have 200 transitions for each of those tasks? ("Context training size in Table 7) This means 4000 transitions total.

However, if we instead sample N=4000 tasks from the semicircle, with one transition each, I would assume that training a classifier becomes problematic. This is what would happen in a continuous task space, and in this sense, the method probably does not work in continuous task spaces.

However, as the authors pointed out in the previous reply, it can generalize from a discrete set of training tasks, so my concerns are resolved about this matter.


I'm also happy with the remainder of the rebuttal and after reading it and the other reviews I will increase my score from 6->7.
Thanks to the authors for their response to my concerns. Regarding Q1, equation 7 should be corrected or clarified concerning the dropping of the $H(M)$ term to avoid confusion. Also, adding the discussion of Q2, Q3, and Q4 in the paper would be nice, as it opens interesting avenues for future work.

Based on the authors' responses and other reviewers' comments, I will keep my scores and increase my confidence.
Thank you for your clarification. There seems to be a misunderstanding about how the context encoder and task representation $Z$ work in UNICORN. Same as FOCAL and other related baselines, our encoder takes a transition tuple $x=(s_t, a_t, r_t, s_{t+1})$ as input and outputs a representation vector $z$ in a $\textit{continuous}$ latent space. Please refer to Figure 6 for visualization of these projected $\textit{continuous}$ representations. At test time, the meta policy performs few-shot generalization by taking new context $X'$ from the test task $M'$, and infer its new representation $Z'$ the same way to condition the downstream RL networks. 

In the example of Ant-Dir, the context encoder generates a $\textit{distinct}$ task representation/statistic for each individual test goal. Therefore, for any goal position on the circle, as long as we are given its context, the context encoder can generate a specific $Z$ for conditioning the RL agent to solve the corresponding task (i.e., navigating to the specific goal position). The specific set of 20 test goals in Figure 1 is only used to empirically demonstrate the generalization of the learned task representation as well as the conditional RL networks. We believe our algorithm should have no problem generalizing to $\textit{any}$ test task sampled from the training distribution (based on result in Sec 3). For out-of-distribution tasks, which is discussed in Figure 1 c) and Sec 4.2, it's indeed extremely challenging, but the generator of UNICORN-SS can still provide nontrivial improvement compared to other baselines, shown in Figure 5.

Again thank you very much for your additional comments. Please let us know if your concern is properly addressed and if there is any misunderstanding or further questions.
Thank you for the rebuttal. I'm a bit busy currently so I will go through the other parts later, but concerning the limitation:

> Q4. discrete task spaces
> The 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.

This seems to be a misunderstanding. I meant that the approach seems to be limited to discrete sets of hidden-parameters, rather than a continuous hidden parameters. I mistyped in my original review, it should be "Adapting them to continuous hidden states should be feasible".

Let's take the examle of Ant-Dir in Figure 1. Intuitively we would like a method that is able to go to any goal position on the circle, not just to the ~20 test goals.
An extension to such a setting would probably be nontrivial and it would be nice to discuss potential ways of doing so in the paper, or at least mentioning this limitation.
[1] http://media.acc.qcc.cuny.edu/faculty/volchok/causalMR/CausalMR3.html

[2] MetaDiffuser: Diffusion Model as Conditional Planner for Offline Meta-RL

[3] ContraBAR: Contrastive Bayes-Adaptive Deep RL

[4] Varibad: A very good method for bayes-adaptive deep rl via meta-learning