Provably Improved Context-Based Offline Meta-RL with Attention and Contrastive Learning
Authors: Lanqing Li
Yuanhao HUANG
Mingzhe Chen
siteng luo
Dijun Luo
Junzhou Huang
Venue: ICLR 2022 Submitted (2022)
Keywords: Reinforcement Learning
Representation learning for planning
Meta-RL
Attention Mechanism
Contrastive Learning
Offline RL

Abstract: Meta-learning for offline reinforcement learning (OMRL) is an understudied problem with tremendous potential impact by enabling RL algorithms in many real-world applications. A popular solution to the problem is to infer task identity as augmented state using a context-based encoder, for which efficient learning of robust task representations remains an open challenge. In this work, we provably improve upon one of the SOTA OMRL algorithms, FOCAL, by incorporating intra-task attention mechanism and inter-task contrastive learning objectives, to robustify task representation learning against sparse reward and distribution shift. Theoretical analysis and experiments are presented to demonstrate the superior performance and robustness of our end-to-end and model-free framework compared to prior algorithms across multiple meta-RL benchmarks.

Under review as a conference paper at ICLR 2022
IMPROVED CONTEXT-BASED OFFLINE META-RL WITH ATTENTION AND CONTRASTIVE LEARNING
PROVABLY
Anonymous authors Paper under double-blind review
ABSTRACT
Meta-learning for offline reinforcement learning (OMRL) an understudied prob- lem with tremendous potential impact by enabling RL algorithms in many real- world applications. A popular solution to the problem is to infer task identity augmented state using context- based encoder, for which efficient learning of robust task representations remains an open challenge. In this work, we provably improve upon one of the SOTA OMRL algorithms, FOCAL, by incorporating intra-task attention mechanism and inter-task contrastive learning objectives, robustify task representation learning against sparse reward and distribution shift. Theoretical analysis and experiments are presented to demonstrate the superior performance and robustness of our end-to-end and model-free framework compared to prior algorithms across multiple meta-RL benchmarks.
INTRODUCTION
Deep reinforcement learning (RL) has achieved many successes with human- or superhuman-level performance across a wide range of complex domains (Mnih et al. 2015) Silver et al., 2017]] Vinyals al., 2019 Ye et al., 2020). However, all these major breakthroughs focus on finding the best- performing strategy by trial-and-eree interactions with single environment, which poses severe constraints for scenarios such as healthcare (Gottesman et al.. 2019). autonomous driving (Shalev Shwartz et al. 2016) and cnntrolled-environe safety is paramount. Moreover, these RL algorithms require tremendous explorations and training samples, and are also prone over-fitting to the target task (Song al. 2019 Whiteson et al. 2011 resulting in poor generalization and robustness. To make RL truly practical in many real-world applications,
agriculture (An et al. 2021 Cao et al. (021) where
paradigm with better safety, sample efficiency and generalization is in need.
Offline meta-RL, as a marriage between offline RL and meta-RL, has emerged as a promising candidate to address the aforementioned challenges. Like supervised learning, offline RL restricts the agent to solely learn from fixed and limited data, circumventing potentially risky explorations. Additionally, offline algorithms are by nature off- off-policy, which by reusing prior experience, have proven to achieve far better sample efficiency than on-policy counterparts (Haarnoja et al.) 2018).
Meta-RL,, the other hand, exploits the shared structure of distribution of tasks and enables the agent to adapt to new tasks with minimal data. One popular approach is by learning a single universal policy conditioned on a latent task representation, known context-based method (Hallak al.)
2015). Alternatively, the shared skills can be learned with a meta-controller (Oh et al. 2017).
In this work we restrict our attention on context-band based offline meta-RL (COMRL), an understudied framework with a few existing algorithms (Li et al. 2019; Dorīman & Tamar 2020 Mitchell et al.
2020 Li et al. 2021a), for set of tasks that differ in reward or transition dynamics. One major challenge associated with this scenario is termed Markov Decision Process (MDP) ambiguity (Li al. 2019) [Dorfman Tamar 2020), namely the task-conditioned policies spuriously correlate task identity with state-action pairs due to biased distribution of the fixed datasets. This phenomenon can be interpreted as a special form of memorization problem in classical meta-learning (Yin et al., 2019). where the value and policy functions overfit the training distributions without capturing causality
Source code provided in the supplementary material.

---

Under review as a conference paper at ICLR 2022
from reward and transition functions, often leading to degenerate task representations (Li al. 2021a) and poor generalization. To alleviate such over-fitting, [i et al. (2021a) proposes a framework named FOCAL which decouples the learning of task inference from control by using self-supervised distance metric learning. However, they made a strong assumption on the existence of an injective map frofrrmm each transistion tuple {s, a, s',r) its task identity. Under extreme scenarios such as sparse reward, where a considerable portion of aggregated experience provides little information regarding task identity, efficient and robust learning of task representations still challenging.
To address the aforementioned problem, in this paper we propose intra-task attention mechanism and inter-task contrastive leamnigg objectives to achieve robust task inference. More specifically, for each task. we apply a batch-wise gated attention to recalibrate the weights of transition samples, and use sequence-wise self-attention (Vaswani et al., 201 to better capture the correlation within the transition (state, action, reward) dimensions. In addition, we implemented a matrix-form objective of the Momentum Contrast (MoCo) (He et al. (2200) for task-level representation learning, by replacing its dictionary queue with meta-batch sampled on-the-fly. We provide theoretical analyses showing that our objective serves as a better surrogate than naive contrastive loss for task inference and the proposed attention mechanism on top can also reduce the variance of task representation. Moreover, empirical evaluations demonstrate that the proposed design choices of attention and contrastive learning mechanisms not only boost the performance of task inference, but also significantly improve its robustness against sparse reward and distribution shift. We name our new method FOCAL++.
RELATED WORK
Attention in RL Although attention mechanism has proven powerful tool across of broad spectrum
of domains (Mnih et al. 2014 Vaswani et al. 2017a Wang & Shen 2017 Veličković et al.) 2018
Devlin al. 2018), to our best knowledge, its applications in RL remain relatively understudied
Most of previous works in RL (Mishra et al.] 2018 Sukhbaatar et al. 2019: Kumar et al. 2020
Parisotto et al. 2020) focus on applying temporal attention order to capture the time-dependent
correlation in MDPs or POMDPs. Raileanu et al. (Raileanu et al., 2020) uses transformer as the
default dynamics/policy encoder for meta-RL., similar to our proposed sequence-wise attention,
without giving any intuition or comparative study on such design choice. |Wang et (202T) recently implemented attention in meta- RL but didn't consider the offline setting.
The closest work we found by far (Barati & Chen 2019 Li et al. (021b) employ attention in multi-viee/multi-agen RI, to learn different weights on various workers or agents, aggregated by a global network to form centralized policy. Analogous our proposal, such architecture has the advantage of adaptively accounting for inhomogencous importance of each input in the decision making process, and makes the global agent robust to noise and partial observability.
Contrastive Learning Contrastive learning (Chopra et al. 2005 Hadsell et al. 2006) has emerged as a powerful framework for representation learning. In essence, it aims to capture data structures by learning to distinguish between semantically similar and dissimilar pairs. Recent progress in contrastive learning focuses mostly on learning visual representations as pretext tasks. MoCo (He) let al. 2020) formulates contrastive learning as dictionary look-up, and builds a dynamic dictionary with a queue and SOTA benchmark with careful composition of data augmentations. However, all these algorithms concentrate primarily on generating pseudo- and contrastive pairs, whereas in COMRL scenario, the task labels and transition samples are naturally given.
avveraee encoder. SimCLR (Chen et al., 2020) further pushes the
There are a few recent works which apply contrastive learning in RL (Laskin et al. 2020 or meta-RU RL (Fu et al. 2020) settings. Fu et al. (2020) employs InfoNCE (Oord et al. 2018) loss to train contrastive context encoder. They investigated the technique the online setting, where the encoder requires an information-gail -giin-baseeddddd based exploration strategy to be effective. In contrast, this paper focuses on how contrastive learning performs in the fully-offline setting.
Based Offline Meta-RL (COMRL) Context-based offline meta-RL employs models with memory such as recurrent Duuan et al. 2016; Wang et al. 2016 Fakoor et al., 2020). recursive (Mishra et al. 2018) probabilistic (Rakelly et al.) 2019) structures achieve fast adaptation by aggregating experience into a latent representation on which the policy is conditioned. To address the bootstrapping error problem (Kumar al. 2019) for offline learning, framework like FOCAL
2

---

Under review as a conference paper at ICLR 2022
enforces behavior regularization (Wu et al. 2019), which constrains the distribution mismatch between the behavior and learning policies in actor-critic objectives. We follow the same paradigm.
METHOD
To tackle the COMRL problem, we follow the procedure described in FOCAL (Li et al. 2021a), by first learning an effective representation of tasks on latent space 2, on which a single universal policy conditioned and trained with bchavior-regularizz actor-critic method (Wu et al. 2019). As an improved version of FOCAL, our main contribution is twofold:
1. To our best knowledge, we are the first to apply attention mechanism in offline multi- task/meta-Rel setting, for learning robust task representations. We combine batch-wise gated attention with sequence-wise transformer encoder, and demonstrate its lower variance as well as robustness against sparse reward and MDP ambiguity compared to prior COMRL methods.
2.
On top of attention, we incorporate a matrix reformulation of Momentum Contrast (He et al. (020) for task representation learning, with theoretical guarantees and provably better performance than ordinary contrastive objective.
3.1 PROBLEM SETUP
Consider family of stationary MDPs defined by (S.S,,,,,RR)) where (S, (S,A,P, R, are the corresponding state space, action space, transition function, reward function and discount factor. A task T is defined an instance of M, which is associated with pair of time-invariant transition and reward functions, P(''ss, a) E P and R(s,a) a) E R, respectively. In this work, we focus on tasks which share the same state and action space. Consequently, a task distribution can be modeled as joint distribution of P and R, which usually can be factorized:
p(T) == RP, p(P).
In the offline setting, each task T; (i being the task label) is associated with a static dataset of transition tuples Dj = { = [(*), 04, *%, [&, a)))). for which p(D,) p(T)). Each tuple e ~ D; is a sequence along the so-called transition/sequence dimension. A meta-batch B is a set of mini-batches
B; ~ D. Consider meta-optimization objective in a multi-task form (Rakelly et al. 2019 Fakoor
et al. 2020
[(0, v) ED, (uuuo (D; ))+ Lerii (en EDD,p(D) [CD, (0, v²),
(2)
where LD, (0, 1/) is the objective evaluated on transition samples drawn from D,, parameterized by and vb Assuming a common uniform distribution for set of n tasks, the meta-training procedure turns into minimizing the average losses across all training tasks
1 - Σ IE [LD, 6]] k=1
i mmota Umeca arg min
(4)
8,0
For COMRL problem, a task distribution corresponds a family of MDPs on which single universal policy is supposed to perform well. Since the MDP family is considered partially observed if no task identity information given, task inference module Eg (* e) is required map context information D to a latent task representation z E Z to form an augmented state, i.c.,
Sang f S Z, Saug concat 2).
(5)
Such an MDP family is formalized as Tusk-Augeented MDP (TA-MDP) in FOCAL. Additionally, Li et al. (2021a) proves that good task representation z is crucial for optimization of the task- conditioned meta-objective in Eqn which is the prime focus of this paper. We now show how to address the issue with the proposed attention architectures and contrastive learning framework.
3.2 ATTENTION ARCHITECTURES
3

---

Under review as a conference paper at ICLR 2022
Task
Roweightng
batch-wise
T1
T1
Atariia
T2
T3
sequence-wise
T3
Attention Block
Attention Block
T
Figure 1: Context encoder as a stack of attention blocks.
We employ two forms of intra-task attention in the context encoder (z batch-wise gated attention and sequence-wiss self-attention, for learning better task representations. The ar- chitectures are shown in Figure 2
Layer Nomr
Layer
Lincar
Multheed Altentinn
Allention
Tanh Activation
Batch- Wise Gated Attention When perform- ing task inference, transitions inside the same batch may contribute differently to the represen- tation learning, especially sparse reward situa- tions. For tasks that differ in rewards, intuitively, transition samples with non-zero rewards con- tain more information regarding the task identity, Therefore, we utilize gating mechanism sim- ilar to (Hu et al. 2018) along the batch dimen- sion to adaptively recalibrates this batch-wise response by computing scalar multiplier for every sample as in Figure [
Layer Nomr
Layer Norm
Sigmold
Postionivise FeedForward
PostfonWss FeedFoowand
AckplivekugPooiid
Rwith Allention
RA Riss Binck
SA Raa Block
Figure 2: Attention modules for task inference. BA: batch-wise attention. SA: sequence-wist wise at- tention.
Sequence-Wis Wise Self-Attention naive MLP encoder maps a concatenated 1-D sequence (s, a, s', from context buffer 1-D embedding z. This seq2seq model can be implemented with sequence-wise attention to apply self-attention along the sequence dimension. The intuition behind sequence-wise attention is that the attentive context encoder should in principle better capture the correlation in (s, s', sequence related to task-specific reward function R(ss, a) and transi- tion function P(sl|s, a), compared normal MLP layers employed by common context-based RL algorithms.
Illustrated in Figure since two attention modules operate on separate dimensions, we connect them in parallel to generate task embedding z by addition.
matrix-form contrastive loss
Figure 3: Inter-task matrix-form momentum contrast. Given two meta-batches of transitions {e*} and {c*) quickly progressing query encoder and slowly progress- ing key encoder compute the corresponding batch-wise wise mean task representations in latent space Z. A matrix mul- tiplication is performed between the set of query and key vectors to produce the supervised contrastive loss in Eqn T.C. Z are the meta-batch, transition and latent space dimensions respectively.
gradient
2 (2*)TT matmul TxT T
T
Z
encoder
(>(e))
momeruum
encoder
(+(e))
T
Tx
4

---

Under review as a conference paper at ICLR 2022
3.3 THE CONTRASTIVE LEARNING FRAMEWORK
Inspired by the successes of contrastive learning in computer vision (He et al., 2020), we process the raw transition data with momentum encoders generate latent query vector 24 as classifier and set of K latent key vectors {=ī. x ...., zx as task representations. Suppose one of the keys = is the only match z", we employ the InfoNCE (Oord al. 2018) objective as the building block:
C. = log exp(=" = /T)
(6)
Σ exp(2" 2// /T)
where 7 is a temperature hyper-parameter (Wu et al. 2018).
To ensure maximum sample efficiency, for each pair of meta-batches B = (B, Díli 1..,, T) where 7' is the meta-batch size, one can construct T' InfoNCE objectives by taking the average latent vector of each task as the key, which also erucial for our theoretical analysis (Theorem 3.1b. Namely, given meta- meta-batch of encoded queries {z] ~ (z 1..... T) and keys
E (z, Bi 1, 1,.., T}, our proposed contrastive loss is
exp(z! 2 =//T) 1ラ=1 exp(2 2// /T)
C,
log
which can be written in a matrix-form
exp(2 2 /T)
= =-T((M Mi log
(8)
Σ/=1 exp(2? 2 /T)
The training scheme of our proposed inter-task momentum contrast is illustrated in Figure 3
Now provide theoretical analysis of the objective in Eqn Intuitively, it is the log loss of T-way softmax-based classifier trying to classify each z as z". With this interpretation, we compare it to a linear classifier with supervised loss and show that it can be recovered by the linear classifier if the weight matrix is a specific mean task classifier (Theorem 3.I). Furthermore, we prove that our proposed objective better surrogate than traditional contrastive loss for task inference.
Definition 3.1 (Supervised Contrastive Loss)
Laup (T (T,g) == EE TT.To ey
where l can be standard hinge or logistic losses in (Saunshi et al. 2019
Consider a linear classifier g(c) = W E(c). where the encoded latent vector E(c) is used as
deterministic representation (Li et al. 202 and W E N is a weight matrix trained minimize Lsup(? (T, W E). Z is the dimension of the task latent space Z. Such construction of contrastive objective enables self- self-supervise task representation learning for task inference, without requiring access to full labels of all possible tasks, which is flexible and has better potential for generalization. Hence the supervised loss of on is defined as
inf 7. (. (T, W W).. T.sup (T, F) WERNxz
(10)
Since the optimal W requires full knowledge and labels of the underlying task distribution, which is infeasible given only training tasks. As with Saunshi et al. (2019). we consider a particular choice of W#:
Definition 3.2 (Mean Task Classifier) For encoder function F. and a task set TT of cardinality N, the mean task classifier W is an N x Z weight matrix whose ith row is the mean latent vector μú of inputs with task label i, We use as shorthand for its loss L' (T
Leup (T, W
In pratice we estimate the mean task representation of 2 and zt using its batch-wise mean
D.D. *
R
GoB,
(11)
(c,)e))
x? E (2dle)
which induces the following definitions:
5

---

Under review as a conference paper at ICLR 2022
Definition 3.3 (Averaged Supervised Contrastive Loss) Average supervised loss for encoder
function E on T-way classification of task representation is defined as
Lusup E)) =
E
(T:) (
(12)
(TAT-1>((T)
The average supervised loss of its mean classifier (Definition 3.2 is
Du (E) =
(TT)EE)(((T) Sup ((Ti)
(13)
When the loss function E is the convex logistic loss, we prove in Appendix B that
Theorem 3.1 The matrix-form momentum contrast objective L2 (Eqn 8 is equivalent to the average supervised loss of its mean classifier Louup (Eqn 13 if E = E (z c) is the key encoder and the mean task classifier WP whose tt row is the mean of latent query vectors with task label i.
If we compare our proposed loss function with the classical unsupervised contrastive loss
Definition 3.4 (Unsupervised Contrastive Loss)
Lun (E) = [(((E())(E(c))>>))) - -(c))))
(14)
Given T as the number distinct tasks in meta-batches, c, c are contexts from the same task, and e is from the other T - 1 tasks. Such construction is employed by prior COMRL methods like FOCAL, which allows for task interpolation during meta-testing.
By Lemma 4.3 in (Saunshi et al. 2019), using convexity of and Jensen's inequality, assuming repeated task labels in each meth-batch, we have
Theorem 3.2 For all context encoder E
Leup ()) S L (E) S (E).
(15)
Combined with Theorem 3.1 it shows that our proposed contrastive objective in Eqn C. =
Li (Eu (Ex(2lc) serves as better surrogate for Lsup than the ordinary unsupervised contrastive losses employed by prior methods, to ensure seallartty-rreervinn task representation for COMRL.
3.4 VARIANCE OF TASK EMBEDDINGS BY FOCAL++
In experiments, we found that our proposed algorithm, FOCAL++, which combines attention mech- anism and matrix-form momentum contrast, exhibit significant smaller variance compared to the baselines on tasks with sparse reward (Table We provide proof of this observation for simplified version of FOCAL++, by only considering the batch-wise attention along with contrastive learning objective defined in Eqn 8, in presence of spurse reward. Assuming all tasks differ only in reward function, we begin with the following definition:
Definition 3.5 (Absolutely Sparse Transition) Given set of tasks [T] which only differ by reward function, a transition tuple (s,a,'', .r) absolutely sparse if VT: E R
constant.
According to policy invariance under reward transformations (gg et al. 1999). without loss of generality, assume the constant above be zero for the rest of the paper.
Definition 3.6 (Task with Sparse Reward) For a dataset D; {(s, (1,,,0,,3 S Re a:))} sampled from any task To with sparse reward, it can be decomposed a disjoint union of two sets of transitions:
D: ((5) ai b/ (s a:)) U (1s,, 04 0),
(16)
= U
(17)
where {c.) } is the set of absolutely sparse transitions (Definition 3.5 which by definition are shared across all tasks. {cn.) consists of the rest of the transitions, and is unique to task Tr.

---

Under review as a conference paper at ICLR 2022
Definition 3.7 (Batch-Wiid Wise Gated Attention) The batch-wise gated attention assigns inhomoge- neous weights W for batch-wise estimation of the mean task representation of μ in Eqn 11
m a (W) EewD, ( (c)
(18)
= W (c.) (c,)) + (c.)
(19)
where Pn. Ps the measures of {cn). {c respectively and W is normalized such that BewD, (c)] = 1. Pm. + Ps by Definition 3.6
Theorem 3.3 Given learned batch-wise gated attention weight W and context encoder E that minimize the contrastive learning objective Leu (W, EE,, we have
Varl (W)) & Varlu
(20)
when the sparsity ratio exceeds threshold.
i.e. the variance learned task embeddings with batch attention upper-bounded by its counterpart without attention given the dataset is sparse enough. We prove Theorem 3.3 in Appendix B.
EXPERIMENTS
In the following experiments, we show FOCAL++ outperforms the existing COMRL algorithms by clear margin in three key aspects: a) asymptotic performance of learned policy; b) task representations with lower variance; and c) robustness to sparse reward and MDP ambiguity.
All trials are averaged over random seeds. The offline training data are generated in accordance with the protocol of FOCAL by training stochastic SAC (Haarnoja et al. 2018) models for every distinct task and roll out policies saved each checkpoint to collect trajectories. The offline training datasets can be collected as selection of the saved trajectories, which facilitates tuning of the performance level and state-action distributions (Table Both training and testing sets are pre-collected, making our method fully-offline Rewards are spursified by constructing a neighborhood of goal in state or velocity space, where transition samples which lie outside the area are assigned zero reward. Since the focus of this paper is robust task representation learning which can be decoupled from control according to FOCAL, we use sparse-reward data only when training the context encoders. Learning of meta-policy in presence of sparse reward is another active but orthogonal area of research where quite a few successful solutions have been found (Andrychowicz et al.) 2017 Eysenbach et al.) 2020)) concrete description of the hyper-parameters and experimental settings is covered in Appendix D.
Table 1: Average testing return (standard deviation parenthesis) of FOCAL and variants of FOCAL++.
Algorithm
Point-Robot Point-Robot-Wind
Sparse-Cheetah-i 1351 10gg
Sparse-Ant-Dir 129 a1.
Sparse-Cheetal--i -183.32/ao.18
Walker-2D-Params
FOCAL
11844m))
p5.61
302, (02.70(22200)
FOCAL+
(contrastvve)
12.33aa)
p.78
1309.76
504.00
-158.95
360.35
FOCAL+
(batch-wise)
12.54
-5.57
1330.56
(87.37(85.00)
150.58
376.52
(seq-wise)
12.64
.5.09
1293 (10.00
573 26
-140.63
375.67(15.7)
TOCAL++
a2.96...o)
5.39 (0.77)
1470.52 (a8.20)
719.77
-137.31
391.02
Table 2: Variance of context embeddings averaged all training tasks and latent dimensions.
Algorithan POCAL
Sparse-Point-Rooo
Point-Robot-We Wind
Sparse-Chectah-Dir Sparse-Ant-Dir
Sparse-Cheetal-Vd
Walker-2D-Params
RS4E-5
3.SEE3
4.11.-
2.24..3
2.57E.3
1.0E-22
FOCAL++
(contrastive)
7.83--
1.68E-3
6.86E-4
1.77E33
1.73E.3
5.79E.3
FOCAL++
(batch-wise)
7.73E-5
1,70E-3
4.66E-4
7.51E-4
1.04E-3
5.85E-3
+0CAL+
(soq-wis)
7.94-55
1.84E-3
9.43E-4
8.00E-4
9.76E-4
5.46E-3
FOCAL+4
8.27E-5
1,68F-3
7.82F.4
1.35E-3
1.06F-3
5.23F-3

---

Under review as a conference paper at ICLR 2022
rece.
Sparsedaint-lthhh lichet
Spars Halk Choetahhhhhhhheee Por ood-Back
1500
130
1000
-250
500
-100
-250
-200
0.4
044
0.5
08
1.
0/
05
1.5
1000 1odtdrendio
Sample Shezs
Serrpe Steps
Sampl Steps
(b) Point-Rohot- Wind
Fwd
Po Ro Wind
320
730
300
200
120
230
-0
150
-9
DV2
0.2
0.4
0.6
0.0
122
0.0
0.4
Sampph Steps
Sarrple Staps
Sampre Steps
POLAL*+
POCAL
BUtch MEAR.
Contentuet BUG
NEM.
(a) FOCAL++ vs. baselines.
(c) Aptrse-Ant-Dir
Figure 4: Left: Test-task performance transition steps sampled for meta-training, Right: t-SNE visualization the learned task embeddings 2 on Point-Robot Wind and Sparse-Ant-Dir. Each point represents query vector which is color-coded according to its task label.
4.1 1 ASYMPTOTIC PERFORMANCE
evaluate +OCAL++ on continuous control meta-cnvironments robotic locomotion (Todorov et al. 2012) adopted from FOCAL. 4 (Sparse-Point PPint-Roboo Sparse-Cheetah- Vel, Sparse-Cheetah- Cheetah-Fwd Back, Sparse- Ant-Fndt--F----- Fadd--d--Bat BackkBkkk and 2 (Point-Robot Wind, Belker--DDPPrammee environments require adaptation by reward and transition functions respectively. For inference, FOCAL+ aggregates context from fixed test set to infer task embedding, and is subsequently evaluated online. Besides FOCAL, three other baselines
compared: an offline variant the PEARL algorithm (Rakelly et al. 2019) (Batch PEARL). a context-based offline BCQ algorithm (Fujimoto et al. [2019) (Contextual BCQ) and two-stage COMRL algorithm with reward/dynamics relabelling (Li et al. 2019) (MBML).
Shown in Figure 4a FOCAL outperforms other methods across almost all domains with context
embeddings of higher quality in Figure 4b 4c In Table [ our ablation studies also show that each design choice of FOCAL++ alone can improve the performance of the learned policy, and combining orthogonal intra-task attention mechanism with inter-task contrastive learning yields the best outcome.
4.2 ROBUSTNESS TO MDP AMBIGUITY AND SPARSE REWARD
In our experimental setup, an ideal context encoder should capture the generalizable information for task inference, namely the difference between reward/dynamics functions across a distribution of tasks. However, as discussed in Section I there are two major challenges that impede conventional COMRL algorithms from learning robust representations:
Table 3: Average testing return of FOCAL and FOCAL+ Sparse-Poine Point-Robor with different distributions of training/testing sets. The numbers in parenthesis represent performance drop due to distribution shift. Additional experiments presented in Apppendix C.
Environment
Training
Testing
FOCAL
FOCAL+
expert
8.16
12.60
expert
medium
7.12(1.04)
12.47
Sparse- Point-
random
4.43(3.13)
10.17
medium
8.44
12.54
Robot
medium
expert
8.25(0.19)
12.44(o.10)
random
6.76
10.49(2.05)
WD.
mixed
302.70
391.02
mixed
Params
expert
271.69
377 377.46(13.56)
8

---

Under review as a conference paper at ICLR 2022
Sparse- Robottrelabeled)
sariples spa samoaks
IOCALar
megosiie Deniity
Auerage See.
04
-1.0
16
0.0
0.2
0.44
0.0
Lest
Batch-arse Vie ghi
(a) State datasct
distribution of relabeled
(b) Test-task performance
(c) Batch-wisc attention weight
Figure 5: Result on the relabeled Sparse-Point-Robee dataset. (a) State distributions of the expert datasets for 20 distinct tasks, with goals uniformly distributed semicircle. (b) On mixed dataset, FOCAL completely fails in this scenario whereas FOCAL++ variants with batch-wise attention are able to learn. (c) Probability distribution of the batch-wise attention weight of samples with absolutely zero and non-zero reward. Binary classification AUC = 0.969.
MDP ambiguity arises due to COMRL algorithms' sensitivity to fixed dataset distributions (Li et al. 2019 Dorfman & Tamar 2020 Take Oparse-Point-Robe for example, as in Figure 5a for tasks with goal on the semicircle, the state-action distribution exhibits specific pattern which may reflect task identity. Given D {s, a, s' r)} as input, the context encoder may learn spurious correlation between state-action distributions and task identity, which causes performance degradation under distribution shifts (Table
Sparse reward in meta-environments could exacerbate MDP ambiguity by making considerable portion of transitions uninformative for task inference, such as the samples outside any goals in Figure 5a Attention mechanism, especially the batch-wise channel attention, helps the context encoder attend to the informative portion of the input transitions, and therefore significantly improve the robustness of the learned policies.
To demonstrate the robustness of FOCAL++ in presence of the two challenges above, we tested against distribution shift by using datasets of various qualities: expert, medium, random and mixed which combines all three. Shown Table 3 we observe that overall the performance drop due to distribution shift is significantly lower when attention and contrastive learning are applied.
Moreover, we are aware that even mixing of datasets generated by different behavior policies cannot fully eliminate the risk of MDP ambiguity since the state-action distributions for each task still do not completely overlap. To show that attention modules introduced by FOCAL++ indeed works intended by capturing the reward-tass dependency, we create new dataset Sparse- Point-Robot by merging the state-action support across all tasks and relabelling the sparse reward according to the task-specific reward functions. In principle, this fully prevents information leakage from the state-action distributions, forcing the context encoder to learn to distinguish the reward functions between tasks while minimizing the contrastive loss. Shown in Figure 5b] experimented with 3 attention variants of FOCAL+ on the relabeled dataset, and found that batch-wise attention significantly improves the performance intended. Additionally, we visualize the density distribution of batch-wise attention weights assigned to samples in Figure 5c, We see clear tendency for the aich-atttetion module assign zero weight to samples with zero rewards (the absolutely sparse data points which lie outside all goal circles in Figure Sap and maximum weights to the non-zero con-zero-reward transitions, with binary classification AUC 0.969, which is clear evidence of FOCAL+ learning the correct correlation for task inference by attending to the informative context.
5 CONCLUSION
In this work, we address the understudied COMRL problem and provably improve upon the existing SOTA baselines such as FOCAL, by focusing on more effective and robust learning of task represen- tations. Key to our framework is the combination of intra-task attention mechanism and inter-task contrastive learning, for which we provide theoretical grounding and experimental evidence on the superiority of our design.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Zhicheng An, Xiaoyan Cao, Yao Yao, Wanpeng Zhang, Lanqing Li, Yue Wang. Shihui Guo, and Dijun Luo. A simulator-based planning framework for optimizing autonomous greenhouse control strategy. In Proceedings of the International Conference on Automated Planning and Scheduling. volume 31, pp. 436 444, 2021.
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay, In Advances in neural information processing systems, pp. 5048-5058, 2017.
Elaheh Barati and Xuewen Chen. An actor-critic-attention mechanism for deep reinforcement learning in multi-view environments. arXiv preprint arXiv: /907.09467. 2019.
Xiaoyan Cao, Yao Yao, Lanqing Li, Wanpeng Zhang, Zhicheng An, Zhong Zhang, Shihui Guo, Li Xiao, Xiaoyu Cao, and Dijun Luo. igrow: A smart agriculture solution to autonomous greenhouse control. arXiv preprint 10t.v:2017.054:: 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Sumit Chopra, Raia Hadsell, and Yann LeCun. Learning a similarity metric discriminatively, with application to face verification. In 2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR'0S), volume 1. pp. 539-546. IEEE, 2005.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint
2018.
Ron Dorfman and Aviv Tamar. Offline meta reinforcement learning, arXiv preprint arXiv:2008.02590 2020.
Yan Duan, John Schulman, Xi Chen, Peter I. Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:1611.02779, 2016.
Benjamin Eysenbach, Xinyang Geng, Sergey Levine, and Ruslan Salakhutdinov. Rewriting history with inverse rl: Hindsight inference for policy improvement. arXiv preprint arXiv:2002.1108) 2020.
Rasool Fakoor, Pratik Chaudhari, Stefano Soatto, and Alexander J. Smola. Meta-g-leamnng. In International Conference on Learning Representations, 2020. URL [attppss//ppeereeiiii net/ /ttp://rrunnndsSeee000E1
Chelsea Finn, Pieter Abbeel, and Sergey Levine Model-agnostic meta-learning for fast adaptation of deep networks. arXiv preprint arXiv: 2017.
Haotian Fu, Hongyao Tang, Jianye Hao, Chen Chen, Xidong Feng, Dong Li, and Wulong Liu. Towards effective context for meea-reinforcemmnt learning: an approach based on contrastive learning. 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, pp. 2052-2062, 2019.
Omer Gottesman, Fredrik Johansson, Matthieu Komorowski, Aldo Faisal, David Sontag, Finale Doshi- Velez, and Leo Anthony Celi. Guidelines for reinforcement learning in healtheare, Nat Med, 26 2019.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. International Conference on Machine Learning, pp. 1861-1870. PMLR, 2018.
Raia Hadsell, Sumit Chopra, and Yann LeCun. Dimensionality reduction by learning invariant mapping. In 2006 IEEÉ Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 06), volume 2, pp. 1735-1742. IEFE, 2006.
10

---

Under review as a conference paper at ICLR 2022
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv preprint arXiv:1502.02259 2015.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9729-9733 2020.
Jie Hu, Shen, and Gang Sun. Squeeze-and- andeezztann-exiaatioon networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7132-7141, 2018.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Advances Neural Information Processing Systems, pp- 11784-11 1794, 2019.
Shakti Kumar, Jerrod Parker, and Panteha Naderian. Adaptive transformers in rl. arXiv preprini arXiv:2004.03761, 2020.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, pp. 5639-5650. PMLR, 2020.
Jiachen Li, Quan Vuong, Shuang Liu, Minghua Liu, Kamil Ciosek, Keith Ross, Henrik Iskov Christensen, and Hao Su. Multi-task Batch Reinforcement Learning with Metric Learning. arXiv e-prints, art. arXiv: 1909. 13373, September 2019.
Lanqing Li, Rui Yang, and Dijun Luo. FOCAL: Efficient fully-offline metu-reinforrcemet learning via distance metric learning and behavior regularization. In International Conference on Learning Representations, 2021a. URL httpss///penrevveww.eet
forum? Bcphh
Wenhao Li. Xiangfeng Wang. Bo Jin, Dijun Luo, and Hongyuan Zha. Structured cooperative reinforcement learning with time-varying composite action space. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021b.
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta-learner. In International Conference on Learning Representations, 2018. URL https
Seppenreview.net/froniee
Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. arXiv preprint arXiv:2008. 06043, 2020.
Volodymyr Mnih, Nicolas Heess, Alex Graves, and Koray Kavukcuoglu. Recurrent models of visual attention. In Proceedings of the 27th International Conference on Neural Information Processing Systems-Volume 2, pp. 2204-2212, 2014.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):5995533, 2015.
Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transformations: Theory and application reward shaping. Ieml, volume 99, pp. 278-287, 1999.
Junhyuk Oh, Satinder Singh, Honglak Lee, and Pushmeet Kohli. Zero-shot task generalization with multi-task deep reinforcement learning. International Conference on Muchine Learning, pp. 2661-2670. PMLR, 2017.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arX 2018.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark. Seb Noury, al. Stabilizing transformers for reinforcement learning. In International Conference on Machine Learning, pp. 7487-7498. PMLR, 2020.
11

---

Under review as a conference paper at ICLR 2022
Roberta Raileanu, Max Goldstein, Arthur Szlam, and Rob Fergus. Fast adaptation to environ- ments via policy-dynamics value functions. In International Conference on Machine Learning, pp. 7920-7931. PMLR, 2020.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off ooflpy meta-reinforcemenn learning via probabilistic context variables. In International conference machine learning, pp. 5331-5340 2019.
Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. Promp: Proximal meta-policy search. arXiv preprini arxiv:/810.06774, 2018.
Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar, A theoretical analysis of contrastive unsupervised representation learning. In International Confer- ence on Machine Learning, pp. 5-28-562266-5 PMLR, 2019.
Shai Shalev-Shwartz, Shaked Shammah, and Amnon Shashua. Safe, multi-agent, reinforcement learning for autonomous driving. arXiv preprint arXiv:1610.03295. 2016.
David Silver, Julian Schrittwieser, Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676)3354-335 359, 2017.
Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational overfitting in reinforcement learning. arXiv preprint 192.0::91..02975, 2019.
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span transformers. arXiv preprint wwXhv:1955.07799, 2019.
Emanuel Todorov, Tom Frez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, pp. 6000-6010, 2017a. URL / 71 B aettett i on i s a 11 yeu-
http:://apest nips
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017b.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference Learning Representations, 2018. URL htttttts://
net wwwp/eeet//eimeeeeeeeeeeemm rJMpiik
Oriol Vinyals, Igor Babuschkin, Wojeiech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung. David Choi, Richard Powell, Timo Ewalds, Petko Georgiev, al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782) 2019.
Jane X Wang. Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint arXiv: 161 005663, 2016.
Jane X Wang, Michael King, Nicolas Pierre Mickael Porcel, Zeb Kurth-Nelson, Tina Zhu, Charlie Deck, Peter Choy, Mary Cassin, Malcolm Reynolds, H Francis Song, et al. Alchemy: A benchmark and analysis toolkit for meta-reiinforemment learning agents. In Thirty-fiith Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Wenguan Wang and Jianbing Shen. Deep visual attention prediction. IEEE Transactions on Image Processing. 27(5):2368-233 2017.
Shimon Whiteson, Brian Tanner, Matthew E Taylor, and Peter Stone. Protecting against evaluation overfitting in empirical reinforcement learning. In 2011 IEEE symposium on adaptive dynamic programming and reinforcement learning (ADPRI.), pp. 120-127. IEER, 2011.
12

---

Under review as a conference paper at ICLR 2022
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arxiv preprint arXiv:1911.1136 2019.
Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via non- parametric instance discrimination. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pp. 3733-3742, 2018.
Deheng Ye, Zhao Liu, Mingfei Sun, Bei Shi, Peilin Zhao, Hao Wu, Hongsheng Yu, Shaojie Yang, Xipeng Wu, Qingwei Guo, al. Mastering complex control in moba games with deep reinforce- ment learning. In AAAI, pp. 6672-6679, 2020.
Mingzhang Yin, George Tucker, Mingyuan Zhou, Sergey Levine, and Chelsea Finn. Meta-learning without memorization. arXiv preprint arXiv:/912.03820, 2019.
13

---

Under review as a conference paper at ICLR 2022
APPENDIX A PSEUDO-CODE
Algorithm I: FOCAL++ Meta-training
Pre-collected batch D; {(8), Og S/. r3) (y))), I3:a....NN from a training tasks (Ti) ==1....... drawn from p(T)
Learning o (22, 03, temperature 7, momentum m
Initialize context replay buffer C, for each task
Initialize context encoder network Eyg (zle), learning policy To (a(ss,,2) and Q-network (8, 2, a) with parameters 2q. Φk, 0 and db
while not done do
for cach T, do
for 1 = T do Sample mini-batches of B transitions {(81,1, & Pailea Dj and update C
end end
Sample pair of query-key meta- bettabbtches of tasks ~ p(T) for step in training steps do for each T, do
11
Sample mini-batches ej and b; ~ C, for context encoder and policy training (b,, are identical by default, the rewards in are always non-sparse) Compute z! Eg)cc)
for each T, do Sample mini- Mini-batches ej from Cj and compute 2) = Eg Mij Mz(2?,
b matrix-form momentum contrast
17
end
L Cactor = Cactor (bs,
18
L Cetitic Coritio (b; 1 (c))
end
21
L2 Tr(M)
22
0g 0g 01 Lz
Φu mΦ + m)
momentum update
0 0 02 Vo E, factor
+ 3 a3 V E, L Coritic
25
end
end
Algorithm 2: FOCAL+ Meta-testing
Pre-collected batch D, = {(s,,, Oj', S'y Ty' (r))rr))) D....... from a set of testing tasks
[Trre=...m drawn from p(T)
Initialize context replay buffer Cv for each task Te
for each Ty do
for 0, T T do
Sample mini-batches of transitions oy Sp Te (e) Dr and update C₁,
Compute 2x.
(cr)
Roll out policy no(als, =8) for evaluation
end
end
To prevent conflict or misunderstanding. all non-hyperlink equatiion/theorem numbers in the appendix refer those in the main
14

---

Under review as a conference paper at ICLR 2022
APPENDIX B DEFINITIONS AND PROOFS
B.1 PROOF OF THEOREM 3.1
Consider a task set (T) = (TI.... Tr) drawn uniformly from p(7). Definition 3.1, the loss incurred by g on point (c, T() E i defined [[((();-- (exs)) which is function of a --ddimensoonnl vector of differences in the coordinates. Given the definition of the mean task classifier W that g(c) W E(c) and f(v) log(1 + Σg exp(-t:)) the standard logistic loss as in (Saunshi et al. 2019). the supervised contrastive loss defined in Eqn 10 can be rewritten as
)
Lisnp (T T,9)= = BTewp(T)
log
1
Σ P'
exp
Σ [We,B(aa)
-
(21)
Since the 8th row of W is the mean of latent key vectors with task label and E (=(cc) is the key encoder, Eqn 21 turns into
Lsup (T,9)== = ETTP((T
log 1+
Σ exp (2 27 x 2
(22)
In practice, we estimate the latent vectors z,th using batch-wise mean to approximate the the mean task representation n0th "hherefore Luup in 22 is equivalent to the mean task classifier L defined
in Definition 3.2. One step futher, assuming uniform distribution of the task set (T) the averaged supervised contrastive loss by Definition 3.3 is
L.up (E) =
E (T.)E
[Lu (1Ti) L)]
log
1 1
exp (2 2 2 29))
(23)
T
1=1
T
exp (2,² x
log
(24)
T
Ej=1 exp
(z 2
1=1
which is precisely the matrix-form momentum contrast objective (Eqn 8,9) if one rescales W by factor T.
B.2 PROOF OF THEOREM 3.3
With Definition 3.5
3.6 and 3.7 we hereby provide a simplified proof by assuming a constant weight
W (c) on the non-sparse } and the absolutely sparse {ca} (Definition 3.5 respectively, then we have
1 (W) = p., -w) (E, (c,, (E.)) + P. W(.) (c.)
E (c,)),
(25)
where the normalization condition E.-D, (c)) = implies pm (c.)) (c..)) + W (cs) = 1. Therefore, adding the batch-wise attention is effectively modulating Ps. and ps- Since Pr. ps 1, without loss of generality, we apply the following notations:
'C {8,,, os, Sy R, (81, a,))} the context space 'Note that the task set {T] discussed here is subset of the whole task set and does not necessarily the whole support of pTT). ). It sampled for the sole purpose computing the contrastive loss.
15

---

Under review as a conference paper at ICLR 2022
Pa (cn) =p Ben ~ feu (cn) xat (arrfc.)EEiii(C =x
(26)
(cn)]
(27)
Assuming i.. In and v,, which gives
Varlu 0i (W) (W)) Var(p' xt (1 - p')==""))) (ar) (1 2 Var(a=?)
(28)
Var(u = Var(px? ?. (1 p)a...1) = p Var(z. + (1 - Var (z
(29)
By B.1 the averaged supervised loss Sup (,, E) is equivalent the matrix-form contrastive objective, which can be written as
L W, E) sup
log
1
Σ r'
exp (1H) (17)) 147)
T
log
1+
Σ V
exp (p' (a
xtt) (4.)
(30)
T
where we use the definition of μ in Equ 25 and the fact that agth k is the same across all tasks. Since (W, (W, F),, and p' ss by the identity map initialization of
the learned W, F. arg min Tu the residual attention module, we have, for learned p', 2 and p.
7 P2P. p.
p <.
(31)
Now subtract Eqn 28 by 29
we have
Var(ng 9.k (W)) Var(u, R.kk = =(17)² - Var [(1 77 -1 -1-p) Var
= V -p) (07 + Var(x:.M)) (2 p--p" Var(
00, is p psg (2 Var(x pVar(a2%.))
(32)
Var(x%** + Var(
The left inequality automatically holds by Eqn 31 the RHS is satisfied when
Var(
p <
(33)
Var(x 2.h
+ Var(
equivalently,
Var(x
P. - pp)
(34)
Var(= 2.
Var(x?
which means when the sparsity of reward exceeds the threshold, learned batch attention module can reduce the variance of the mean task representation n g. Eqn 34 is corroborated by our experiments on the relabeled Sparse-Poin Point- Robot dataset (Figure 6).
16

---

Under review as a conference paper at ICLR 2022
APPENDIX C ADDITIONAL EXPERIMENTS
In Table we present more experimental evidence that +OCAL+ more robust against distribution shift compared to FOCAL on Walker-2D-Params, which is consistent with Table 3 in the main text.
Table 4: Extension of Table in the main text. Average testing return of FOCAL and FOCAL++ for more settings of distribution shift on Walker-2D-Parame
Environment
Training
Testing
FOCAL
FOCAL++
expert
373.92
364.75
expert
mixed
322.24 68)
340.60(21.15)
Walker-2D- Params
random
284.94 98)
297. 297.43
mixed
302.70 70
391.02
mixed
expert
271.69(31.0) (1)00)
377.46 37T.46(13.56)
random
260.02 (42.66))4..)
3-16.95(44.07)
1
ROCAL FOCAL+
25-
variance
1.5
0.1
0.5
0.6
0.7
0.8
0.9
1.0
sparse ratio
Figure 6: The variance-sparsity relation for FOCAL++/FOCAL the relabeled Yparse-Point-Roboe dataset. The y-axis measures the variance of the bounded task embeddings E (-1, 1)' averaged over all latent dimensions. See more details in D.2
Moreover, to testify our conclusion in B.2 present the variance of task embedding vectors of FOCAL++ and FOCAL under various sparsity levels. Shown in Figure the variance of the weighted embeddings 1 (W) becomes lower than its unweighted counterpart μa q.k wwen sparse ratio exceeds threshold about 0.6. The observation matches well with Eqn 34 we derived in B.2
17

---

Under review as a conference paper at ICLR 2022
APPENDIX D EXPERIMENTAL DETAILS AND HYPERPARAMETER
D.1 OVERVIEW OF THE META ENVIRONMENTS
The meta-environmentt could be divided into two categories: meta-envviromment that only differ in reward function and that only differ in transition function. For the eevareenvioomments that only differ in reward functions, additionally introduce sparsity to the reward function.
Sparse-Point-Roiin is 2-navigation task with sparse reward, introduced in Rakelly et al. (2019). Each task associated with goal sampled uniformly on unit semicircle. The agent is trained to navigate to set of goals, then tested on distinct set of unseen test goals. Tasks differ in reward function only.
Point-Robot- Wind is another variant of Sparse-Point-Rooee Each task is associated with the same reward but a distinct "wind" sampled uniformly from [-1,0]. Every time the agent takes step, it drifts by the wind vector. We set 0.05 in this paper. Tasks differ in transition function only.
Crrrse-Cheetah-Vee Vel, Dparse-Ant-Fwd-Bacck reward variants of the popular meta-RL benchmarks and Saarse-Chhetah-Fwwl Brcck based on MuJoCo environments, introduced by Finn et al. (017)) and Rothfuss et (2018). Tasks differ in reward function only.
Hearse-Cheetah-Fwweee
are sparse- Vel, Sparse-Ant-Dae
Walker-2D-Parame is unique environment compared to other MuloCo environments. Agent is initialized with some system dynamics parameters randomized and must move forward. Transitions function is dependent on randomized task-specific parameters such as mass, inertia and friction coefficients. Tasks differ in transition function only.
The way sparsify the reward functions is follows.
reward - goal radius goal radius
if reward > goal radius
sparsified reward
(35)
0,
otherwise
Intuitively, we set rewards of states that lie outside neighborhood of the goal to 0, and re-scaled the rewards otherwise that the sparse reward function continuous. For each of the sparsified Cearse-Point-Ronee
environments other than the relabeled sparse rate of about 50%. Note that only the transitions used for training the centext-encoder are sparsified, since the focus of this paper is learning effective and robust task representations.
we set its goal radius to achieve non-
D.2 RELABELED DATASET
As discussed in Section 4.3, to prevent information leakage task identity from state-action distri- bution, construct the relabeled Sparse-Point-Rohot dataset from a pre-collected dataset of the Sparse-Point-Ront environment.
Figure 7 illustrates the generating process for task 2 of the original dataset. The original state distribution of five example tasks on Sparse-Point-Robes is shown in the upper-left. After merging the transition state-action support across all tasks, the (state, action, next state) distribution are identical for every specific task. Then we recompute the reward for each transition according to the task specific reward functions and sparsify the result. We perform the merge-relabel-sparsii process for all tasks on Sparse- Point-Poioooooeeee to enhance the importance of the non-sparse samples for task inference. The sparse samples Figure 7 of the main text are those that lie outside of all goals, i.e. transitions with zero reward across all tasks.
The dataset can be accessed and downloaded from relabeled dataset
D.3 HYPERPARAMETTRS
Tables 5 and 6 describe the hyperparameters used in empirical evaluations.
18

---

Under review as a conference paper at ICLR 2022
Goal 4
Goal
Merge transitions
5
Goal
-03
Relabel rewards
Sparsify
Goal
rewards
Goal
pawe
455
Figure 7: Generating process the relabeled Sparse-Point Poiee-PPontt-ooooo Robot dataset.
Table 5: Specifications of the environments experimented in our paper.
Training Set
Training Tasks
Testing Tasks
Goal Radius
Sparse Pacee-Poin---eee Point-Robot
80
-0.2
Sparse-Point Point-Robo (relabeled)
Point- Roaett-oooo---o Robot-Wind
40
Sparse- Crrise-Chhetah-WW Vel
80
Sparse-Ant Fwd--Bakk
2
2
Walker2d Rand-Params
20
N/A
D.4 IMPLEMENTATION
All experiments are carried out 64-bit CentOS 7.2 with Tesla P40 GPUs. Code is implemented and run with PyTorch 1.2.0. One can refer to the source code in the supplementary material for complete list of dependencies of the running environment.
19

---

Under review as a conference paper at ICLR 2022
Table 6: Hyperparameters used for training to produce Figure 4(a). Meta-batch size refers to the number of distinct tasks for computing the DML contrastive loss at time. Larger meta-batee size leads to faster convergence but requires greater computational power. For Fwd-Back environments, meta-batch size of suffices for stability and efficiency.
Hyperparameters reward scale discount factor
Point-Robor
Mujoco
100
5
0.9
0.99
maximum episode length
20
200
target divergence
N/A
0.05
behavior regularization strength(a)
500
latent space dimension
20
meta-batch size dml_Ir(a;)
16*
le-3
3e-3
actor_Ir(a2)
le-3
3e-3
critic_ir(ag)
le-3
3e-3
DML. loss weight(3) contrastive T
0.5
0.5
contrastive m buffer size (per task)
0.9
le4
1e4
batch size (sac)
256
256
batch size (context encoder) g If(f--ivveeeence discriminator)
512
512
le-4
le-4
transformer hidden size (context encoder)
128
128
multihead (if enabled)
8
reduction (batch attention)
16
16
transformer blocks (context encoder)
3
dropout (context encoder) network width (others)
0.1
0.1
256
256
network depth (others)
3
20


Reviews:
[{"review_id": "xXuN9QBObqK", "reviewer": "Reviewer_UknQ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper makes a valuable and novel contribution in the sense that offline meta-RL is somewhat under-studied and needs better methods. I also appreciated the focused study on context-based methods, and it seems like the proposed modifications could be widely applicable. However, I have some concerns on experimental and theoretical results.\n\n**Experiments**\n\nThe experimental design leave much to be desired.\n\n1. I do not understand the logic behind \"using sparse-reward data only to learn context encoders\": I assume the policy is learned with dense offline rewards, which seems unrealistic — either the observed rewards are dense or they are sparse. This seems to stem from assuming that encoder and policy are decoupled, but this is likely false: the policy is trained given a particular encoder, and replacing the encoder with a different one will degrade performance. In any case, I wished the paper included results on the standard non-sparse environments.\n2. Similarly, I missed common baselines in offline meta-RL, e.g., MACAW (Mitchell et al., 2021) or BOReL (Dorfman et al., 2020). This is an issue because it's unclear how strong the included baselines (not specifically designed for offline RL) are against those methods, and thus limits potential interest from the community. If instead FOCAL++ was competitive with or outperformed those methods on (sparse & dense reward) tasks, this would be a strong selling point for the paper.\n3. It seems like the proposed components (Attention & MoCo) could be combined with other context-based offline meta-RL methods (e.g., offline RL^2). Is it the case? If yes, showing such results would significantly strengthen the submission.\n4. On the positive side, I appreciated the discussion on MDP ambiguity and sparse rewards (p. 9), which emphasize the importance of Attention (and MoCo) when task identification is challenging.\n\n**Theory**\n\n1. I find the results of theorems 3.1 and 3.2 misguiding. Th. 3.1 is just a matter of renaming the i-th column of $W^\\mu$ as $\\mathbb{E}_k[z^k_i]$, and I don't see a contribution there. As I understand it, Th. 3.2 simply invokes the results of Saunchi et al., 2019, so it seems the contribution is also very minor.\n2. Th. 3.3 is more interesting, especially for context-based methods in sparse-reward environments.\n\n**Minor details**:\n\n- Recent work has already used attention for meta-RL (c.f., claim in 1st bullet point of p. 3): for example, Wang et al., 2021 use a transformer-based policy to meta-learn on tasks in their proposed benchmark.\n- On p. 3, before Eq. 1: \"which usually can be factorized\" is missing \"which\"?\n\n**References**\n\n- Mitchell et al., \"Offline Meta-Reinforcement Learning with Advantage Weighting\", ICML 2021.\n- Dorfman et al., \"Offline Meta Learning of Exploration\", ArXiv 2020.\n- Wang et al., \"Alchemy: A benchmark and analysis toolkit for meta-reinforcement learning agents\", ArXiv 2021.", "summary_of_the_paper": "This submission tackles the problem of offline meta-RL, where the goal is to learn a policy that can quickly solve new tasks. This policy must be learned on an offline set of (pre-collected) trajectories from multiple tasks, unlike in online meta-RL where data collection and learning are interleaved.\n\nSpecifically, the authors focus on improving a previous method, FOCAL, in 2 aspects:\n\n1. They add an attention mechanism to improve the estimation of the context vector identifying the task. This mechanism attends to both the transitions from one state (and rewards, actions) to the next within a trajectory, and across trajectories within a mini-batch from a single task. The authors show that this mechanism can reduce variance when estimating task context.\n2. They propose to use MoCo (He et al., 2020) to improve and accelerate learning of the task encoder. They include theoretical results showing that MoCo better approximates a supervised objective.\n\nThey conclude with experimental results on point-mass and simple Mujoco tasks, where FOCAL++ compares favorably against FOCAL and contextual variants of meta-RL algorithms.", "summary_of_the_review": "Overall, I think this paper proposes reasonable modifications for context-based offline meta-RL methods. But, I worry that their results focus too narrowly on FOCAL to be of wide interest to the ICLR community.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "etuGkLe-xA-", "reviewer": "Reviewer_uEZk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n1. The paper focuses on a very important problem in reinforcement learning: meta-learning. The ability to adapt to new tasks and\ntransfer learning is extremely important component of intelligence.\n\n2. The authors firmly ground their approach with mathematical motivation and proofs. \n\n3. The quantitative results on the given Mujoco dataset are competitive and show promise.\n\nWeaknesses:\n\n1. The conceptual novelty is limited. The paper effectivly just takes FOCAL and adds an attention mechanism and a contrastive loss. \n\n2. The approach was evaluated on only one dataset. It is difficult to draw conclusions on the generality of the approach. An evaluation \non another meta-learning environment (such as perhaps Meta-World) would strengthen the case of the paper.", "summary_of_the_paper": "In this paper, the authors focus on the problem of meta-learning for offline reinforcement learning. To this end, they build \non the FOCAL algorithm by adding to key contributions. First, they add an intra-task attention mechanism. Second, they add\ninter-task contrastive learning. The authors provide a theoretical foundation for their approach and evaluate on control tasks\nin MuJoco. They demonstrate competitive performance over the state of the art.", "summary_of_the_review": "In general, the authors have proposed a modification to an existing algorithm, FOCAL, which incorporates attention and contrastive learning.\nThey evaluate on only one dataset. In spite of these weaknesses, the approach has strong mathematical foundations, and results on the Mujoco dataset look promising. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "XU5QORwsurq", "reviewer": "Reviewer_oBW4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper provides a descent improvement for the SOTA offline meta RL algorithm, FOCAL. The main contribution is a combination of three existing techniques that are widely used in other domains, and therefore the novelty / significance is somewhat limited. The experiment results do show the improved test returns and reduced variance in all the six tasks, and the ablation studies are helpful to understand the benefit of each component and its advantage under distribution shift. Therefore, this work is marginally above the acceptance threshold.\n\nThe theoretical analysis is somewhat interesting, but I am not sure how much value it adds to the paper and the proof in Theorem 3.3 is not solid.\n\nFirst, the equivalence between the matrix-form momentum contrast objective and the mean classifier loss is quite straightforward by observing Eq 8 as the log-soft-max loss with the logits being the linear product of latent vectors. The main theorem 3.2 in Section 3.3 is an application of Lemma 4.3, but I don’t find any details about how it is applied. While we can say L_{sup}^\\mu (13) is closer to the supervised contrastive loss than the unsupervised contrastive loss (14), it is hard to argue that the former serves as a better surrogate simply because of that, neither (13) will learn a better task embedding than (14).\n\nSecond, I’m not sure if theorem 3.3 is rigorously proven in the appendix. A key step of the proof is Eq 31, but the argument in the text between 30 and 31 is rather vague. I do not see how 31 is derived. Even if Theorem 3.3 is true, it only shows that when the context encoder with a batch-wise gated attention is optimised, the variance of the embedding after the attention is smaller than the variance of equally weighted embedding before that layer. But it does not show whether it is smaller than the variance of the learned embedding if one optimises the original context encoder without the attention.\n\nIn the experiment section, it seems the contribution of the contrastive loss is only marginal from Table 1 and 2. How would FOCAL++ perform if we only use the batch-wise and seq-wise attention? Also, I suppose adding the two attention architecture increase the overall size of the context encoder. Will we improve the performance of FOCAL with a similar gain by using a bigger network with the original architecture?\n", "summary_of_the_paper": "This paper proposes to improve a context-based offline meta-RL algorithm, FOCAL, by several modifications to the context encoder’s architecture and training objective. Specifically, it adds batch-wise gated attention and sequence-wise self-attention in the encoder network, and proposes to replace the classical unsupervised contrastive loss with the InfoNCE objective. Theoretical analysis is done to show the proposed objective is closer to the supervised contrastive loss and the variance of task embedding with an optimised encoder is smaller than the one without batch-wise gated attention. Experiments on a few common meta RL benchmarks and sparsified environments show the superiority of the method over FOCAL and a few other baselines.", "summary_of_the_review": "The architecture improvement on FOCAL is a good contribution in practice, while the theoretical analysis part requires some work for more solid discussion and proof.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "Mneavn20PQo", "reviewer": "Reviewer_yZzz", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- Description of the method is reasonably easy to follow\n- Interesting experiments on distribution shift and MDP ambiguity\n\nWeaknesses:\n- The scope of the paper is rather narrow (specifically context-based offline meta-RL, mostly focused on improving a single algorithm), which will probably limit impact; consider adding comparisons with [1] or [3]?\n- The technical novelty is limited (essentially, adding attention)\n- It's very difficult to interpret the significance of the main theoretical claim; can the authors justify why L_{sup} is a meaningful lower bound? Where does this loss come from? Unless I'm missing something, the solution of W = 0 would give zero loss, depending on the choice of \\ell\n- The writing gets very difficult to follow in the theoretical sections; variables are not always defined unambiguously and some theoretical statements lack justification/explanation\n\nOther comments:\n- The term \"MDP ambiguity\" doesn't come from [2], but [1] (as far as I know)\n- \"The intuition behind sequence-wise attention is that the attentive context encoder should in principle better capture the correlation in (s, a, s′, r) sequence related to task-specific reward function R(s, a) and transition function P (s′ |s, a), compared to normal MLP layers employed by common context-based RL algorithms.\" MLP layers are already universal function approximators, so it's not clear to me why self-attention is needed to learn a better representation of the tuple. If the tuple. Self-attention seems useful if the sequence is too long to fit into a regular MLP or if the sequence length is variable, but neither of these are the case here, as far as I can tell.\n- It seems to me that the matrix form of the contrastive loss is an implementation detail, rather than a novelty, but maybe I'm missing something...\n- The notation in section 3.2 and 3.3 is a bit confusing, because the same symbol (e.g. bold z) means a matrix in 3.2, but a single vector in 3.3.\n- Without knowing what \\ell is in Eq 9, it's hard to know what the equation means. Currently, it's not clear that this \"supervised contrastive loss\" is a meaningful objective. Assuming that \\ell(0) = 0, doesn't the zero matrix give zero loss?\n- I don't really follow the interpretation of the key vectors as classifiers- wouldn't the query classify whether or not particular keys are of the same class?\n- Definition 3.1: It's unintuitive to me to define the loss for a classifier as the difference between their outputs, as opposed to e.g. KL divergence or cross entropy. Why was this form chosen?\n- Are the queries formed by simple averaging, or the weighted average from self-attention?\n- \"Since we assume no access to the entire task set\" This is strange to me- we're using the task set during training, right?\n- The writing in Sec 3.3 is difficult to follow. In particular, I didn't realize that we were trying to come up with a surrogate objective for the \"supervised loss\" until the middle of page 6. It's also not clear to me that the supervised loss in Eq 9 is a meaningful thing to bound.\n- if the constant is assumed to be zero, I'm not sure definition 3.5 makes sense- doesn't this mean all transitions have reward zero?\n- Section 3.4 says that a proof for Theorem 3.3 is given in the Appendix, but the Appendix lists an \"informal proof\", which is a bit strange\n\n[1] Dorfman & Tamar. Offline Meta-Learning of Exploration. 2020.\n[2] Li et al. Multi-task Batch Reinforcement Learning with Metric Learning. 2020.\n[3] Mitchell et al. Offline Meta-Reinforcement Learning with Advantage Weighting. 2021.", "summary_of_the_paper": "The paper intends to improve a recent offline meta-RL algorithm (FOCAL) through a modified objective and architecture. The authors propose the use of attention both for up-weighting more informative transition tuples in a batch of adaptation data, as well as within the transition tuples themselves, in order to improve task inference in sparse reward settings. In addition, the authors modify the objective used in FOCAL to form a tighter bound on a \"supervised\" contrastive loss than prior work, though it's not totally clear how meaningful this supervised form of the loss is.", "summary_of_the_review": "I have mixed feelings on this paper. Overall, I think the technical novelty and impact of theoretical contributions is potentially limited, but some of the experiments in distribution shift/MDP ambiguity are interesting. Currently I'm not comfortable accepting the paper on account of dubious theoretical contributions and limited scope/technical novelty, but if the authors can convincingly argue why their bound in Theorem 3.1 is meaningful (motivate the supervised version of the loss), add additional baselines to make the paper less specific to just FOCAL/COMRL, and improve the writing in the theoretical sections, I am open to accepting it (but could be convinced otherwise).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper addresses the problem of offline meta reinforcement learning. The authors build on the FOCAL algorithm, adding intra-task attention and inter-task contrastive representation learning objectives. The resulting FOCAL++ algorithm outperforms several strong baseline, including FOCAL and a theoretical analysis attempting to show that FOCAL++ provably improves on FOCAL is included.

Reviewers agreed that the novelty of the proposed approach is limited since attention and contrastive representation learning have been used in the closely related (online) meta-RL setting. At the same time reviewers agreed that the results in the paper and the rebuttal show that FOCAL++ improves on a strong set of baselines.

The main shared concern was regarding the significance and validity of the theoretical analysis. After considering the rebuttal reviewers voting for both acceptance and rejection were in agreement that there are issues with the theoretical analysis/justification. While we agree with the authors that the algorithmic and experimental part of the paper is strong, we have to base our decision on the state of the whole paper. In the end we decided not to accept the paper because 1) the paper put a significant focus on a theoretical argument the reviewers found problematic and 2) the authors did not modify the paper to sufficiently address these concerns during the available window.
Decision: Reject