We sincerely thank the reviewer for the constructive comments, which really help us make the paper stronger. Our response to your concerns/questions:

Justification of Assumption 2.2

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

In 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].

Hyper-Parameters

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

For all methods on each environment, we do thorough grid-search to select the hyper-parameters:

Prompt-DT: $K~[2,5,10,40]$, same as its paper.

MACAW: $\lambda\sim [0.005, 0.01, 0.05, 0.1]$, where 0.01 is recommended in its paper.

FOCAL: $\beta\sim[1,2,8,16]$ and $n\sim[-2,-1,1,2]$, same as its paper.

CORRO: we fine-tune the training process of the negative-generator, reducing the training loss from 2.833 (open-source repo) to 0.084.

CSRO: $\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.

UNICORN-SUP: no hyper-parameter to tune. Hence, UNICORN-SUP is recommended if no hyper-parameter tuning is a priority.

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

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

Thm 2.4

Please refer to our following response to Q3.

How UNICORN-SUP encodes the context

We use the output feature of the encoder as $z$, which is also the input of the classifier.

Questions

Q1. "Interpretation of $I(Z; M)\ge I(Z; M|X_b)$"

We now provide a rigorous proof:

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

$\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$

$\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$ . 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) \ge const$

Namely, 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.

Q2. ContraBAR

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

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

Q3. Thm 2.4

Please see our general response Q4.

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.

[1] http://media.acc.qcc.cuny.edu/faculty/volchok/causalMR/CausalMR3.html

[2] Conditional Mutual Information to Approximate Causality for Multivariate Physiological Time Series

[3] On Causality and Mutual Information

[4] Soft actor-critic for discrete action settings

[5] Human-level control through deep reinforcement learning

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.

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 $*continuous*$ latent space. Please refer to Figure 6 for visualization of these projected $*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 $*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 $*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 comments. The "context training size" refers to the number of transitions used for each task during $*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?

We sincerely thank the reviewer for the constructive comments and the appreciation of our work. Our response to your concerns/questions:

Q1. Eqn 7

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

Q2. "How is convex interpolation beneficial?"

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

Q3. Scope of the framework

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

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 [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).

Q4. Extending UNICORN to Online Meta-RL

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

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

[2] ContraBAR: Contrastive Bayes-Adaptive Deep RL

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

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.

We sincerely thank the reviewer for the constructive comments and the appreciation of our work. Our response to your concerns/questions:

Clarity

Thank you for your advice. We will make sure to improve the presentation of figures and theorems in the final version.

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

Adding $I(Z; X) = I(Z; X_t|X_b) + I(Z; X_b)$ in Theorem 2.3 is a good idea.

Significance

• proof sketch:

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

• why UNICORN outperforms baselines:

As explained in our following reponse to Q2, the lesser causality $*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 $*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.

Questions

Q1. "why would the red line in Fig 1a) not be more correlated with the goal direction?"

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

Q2. "why including any lesser causality would improve the latent representations?"

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

Q3. "Transition" vs "Trajectory"

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

Q4. Last inequality of eqn 14.

Since 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)$

Proof: let $q$ denote $p_{\theta}$ in eqn 14,

$-I(X_t; Z, X_b) \equiv -\mathbb{E}_{x_t, x_b, z} [\log p(x_t|z, x_b)]$

$= -\mathbb{E}_{x_t, x_b, z}[\log \frac{p(x_t|z, x_b)}{q(x_t|z, x_b)}]$

$-\mathbb{E}_{x_t, x_b, z}[\log q(x_t|z, x_b)]$

$= -\mathbb{E}[KL(p(x_t|z, x_b)||q(x_t|z, x_b))] + L_{\text{recon}} \le L_{\text{recon}}$

Q5. "how the context is obtained from the replay buffer"

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

For downstream RL training: we randomly sample a batch of single-step transitions $x$ from the replay buffer for policy optimization.

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

Q6. "it should be impossible to infer task from random-quality dataset"

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

[1] Representation learning with contrastive predictive coding

[2] Provably improved context-based offline meta-rl with attention and contrastive learning

[3] Auto-encoding variational bayes