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

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

Strengths

* The information theoretical analysis seems novel, especially in its analysis and comparison of the prior existing methods. * The new proposed methods are well motivated based on the prior analysis and perform well on a diverse set of problem setting variations * 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. * 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

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

Rating

7

Confidence

3

Soundness

2

Presentation

2

Contribution

3

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.

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.

Strengths

- This paper is very well motivated by highlighting previous work contributions to the field and how they fit with their framework. - 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. - 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. - 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). - 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

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.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors discuss some limitations, and the questions above may address potentially new ones.

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.

Strengths

# Originality The 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. # Clarity Overall 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. # Quality The 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. # Significance The 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. # 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

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?

Rating

8

Confidence

4

Soundness

3

Presentation

4

Contribution

4

Limitations

Limitations are stated and clear.

Comment by Reviewer sHdF

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.