Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning
Authors: Rishabh Agarwal
Marlos C. Machado
Pablo Samuel Castro
Marc G Bellemare
Venue: ICLR 2021 Spotlight (2021)
Keywords: Reinforcement
Generalization
Contrastive learning
Bisimulation
Representation Learning

Abstract: Reinforcement learning methods trained on few environments rarely learn policies that generalize to unseen environments. To improve generalization, we incorporate the inherent sequential structure in reinforcement learning into the representation learning process. This approach is orthogonal to recent approaches, which rarely exploit this structure explicitly. Specifically, we introduce a theoretically motivated policy similarity metric (PSM) for measuring behavioral similarity between states. PSM assigns high similarity to states for which the optimal policies in those states as well as in future states are similar. We also present a contrastive representation learning procedure to embed any state similarity metric, which we instantiate with PSM to obtain policy similarity embeddings (PSEs). We demonstrate that PSEs improve generalization on diverse benchmarks, including LQR with spurious correlations, a jumping task from pixels, and Distracting DM Control Suite. 

Published as a conference paper at ICLR 2021
CONTRASTIVE BEHAVIORAL SIMILARITY EMBEDDINGS FOR GENERALIZATION IN REINFORCEMENT LEARNING
Rishabh Agarwal' Marlos C. Machado
Pablo Samuel Castro
Mare G. Bellemare
Google Research, Brain Team
srishabhagarwal, marlosm, psc, be bellemare ) @google. com
ABSTRACT
Reinforcement learning methods trained on few environments rarely learn policies that generalize to unseen enviromments. To improve generalization, we incorporate the inherent sequential structure reinforcement learning into the representation learning process. This approach is orthogonal to recent approaches, which rarely exploit this structure explicitly. Specifically, we introduce a theoretically motivated policy similarity metric (PSM) for measuring behavioral similarity between states. PSM assigns high similarity to states for which the optimal policies in those states well as in future states are similar. We also present contrastive representation learning procedure to embed any state similarity metric, which we instantiate with PSM obtain policy similarity embeddings (PSEs)). We demonstrate that PSEs improve generalization on diverse benchmarks, including LQR with spurious correlations, jumping task from pixels, and Distracting DM Control Suite.
INTRODUCTION
TRAIN ENV
TRAIN ENV
TEST ENV
Current reinforcement learning (RL) approaches of- ten learn policies that do not generalize to environ- ments different than those the agent was trained on,
even when these environments are semantically equiv- alent (Tachet des Combes et al., 201 8; Song et al., 2019; Cobbe al., 2019). For example, consider a jumping task where an agent, learning from pixels, needs to jump over an obstacle (Figure 1). Deep RL agents trained on few of these tasks with different obstacle positions struggle to solve test tasks where obstacles are at previously unseen locations.
The agent (white block), learning from pixels, needs to jump over an obsta- cle (grey square). The challenge is to generalize to unseen obstacle positions and floor heights in test tasks using small number of training tasks. We show agent's 's trajectories using faded blocks.
Figure I:
Jumping task:
Recent solutions to circumvent poor generalization in RL are adapted from supervised learning, and, as such, largely ignore the sequential aspect of RL. Most of these solutions revolve around enhancing the learning process, including data augmentation (e.g., Kostrikov et al., 2020; Lee et al., 2020a), regularization (Cobbe et al., 2019; Farebrother et al., 2018), noise injection (Igl et al., 2019), and diverse training conditions (Tobin et al., 2017); they rarely exploit properties of the sequential decision making problem such as similarity in actions across temporal observations.
Instead, we tackle generalization by incorporating properties of the RL problem into the representation learning process. Our approach exploits the fact that an agent, when operating in environments with similar underlying mechanics, exhibits at least short sequences of behaviors that are similar across these environments. Concretely, the agent is optimized to learn an embedding in which states close when the agent's ' optimal policies in these states and future states are similar. This notion of proximity is general and is applicable observations from different environments.
Specifically, inspired by bisimulation metrics (Castro, 2020; Ferns et al., 2004), we propose a novel policy similarity metric (PSM). PSM (Section 3) defines notion of similarity between states originated from different environments by the proximity of the fong-term optimal behavior from these states. PSM is reward-agnostic, making it more robust for generalization compared to approaches that
Also at Mila, Université de Montréal. Now at DeepMind.
Pronounce "pisces'.

---

Published as a conference paper at ICLR 2021
rely on reward information. We prove that PSM yields an upper bound on suboptimality of policies transferred from one environment to another (Theorem 1), which is not attainable with bisimulation.
We employ PSM for representation learning and introduce policy similarity embeddings (PSEs) for deep RL. To do so, we present a general contrastive procedure (Section 4) to learn an embedding based any state similarity metric. PSEs are the instantiation of this procedure with PSM. PSEs are appealing for generalization as they encode task-relevant invariances by putting behaviorally equivalent states together. This is unlike prior approaches, which rely on capturing such invariances without being explicilly trained to do so, for example, through value function similarities across states (e.g., Castro & Precup. 2010), or being robust to fixed transformations of the observation space (e.g.. Kostrikov et al., 2020; Laskin et al., 2020a).
PSEs lead better generalization while being orthogonal to how most of the field has been tackling generalization. We illustrate the efficacy and broad applicability of our approach on three existing benchmarks specifically designed to test generalization: (i) jumping task from pixels (Tachet des Combes et al., 2018) (Section 5), (ii) LQR with spurious correlations (Song et al., 2019) (Section 6.1), and (iii) Distracting DM Control Suite (Stone et al., 2021; Zhang et al., 2018b) (Section 6.2). Our approach improves generalization compared to a wide variety of approaches including standard regularization (Farebrother et al., 2018; Cobbe et al., 2019), bisimulation (Castro & Precup, 2010; Castro, 2020; Zhang et al., 2021), out-of-distribution generalization (Arjovsky et al., 2019) and state-of-tth-arr data augmentation (Kostrikov et al., 2020; Lee et al., 2020a).
2 PRELIMINARIES
We describe an environment as a Markov decision process (MDP) (Puterman, 1994) M A, P, 7), with a state space X, an action space A, a reward function R. transition dynamics P. and discount factor ? [0, 1). A policy T(- x)) maps states x E X' distributions over actions. Whenever convenient, we abuse notation and write (z) to describe the probability distribution at zz). treating 7 (z)) as a vector. In RL, the goal to find an optimal policy ' that maximizes the cumulative expected return Eu,~
mx) |E. 7 R(Ia, 4t)] starting from an initial state zo-
We are interested in learning a policy that generalizes across related environments. We formalize this by considering a collection p of MDPs, sharing an action space A but with disjoint state spaces. We use X' and y to denote the state spaces of specific environments, and write Rx, Px for the reward and transition functions of the MDP whose state space is X, and xx for its optimal policy, which we assume unique without loss of generality. For a given policy π, we further specialize these into Rx and PT, the reward and state-to-state transition dynamics arising from following π in that MDP.
We write for the union of the state spaces of the MDPs in P. Concretely, different MDPs correspond to specific scenarios in problem class (Figure 1), and S is the space of all possible configurations. Used without subscripts, R. P. and π refer to the reward and transition function of this "union MDP", and policy defined over S; this notation simplifies the exposition. We measure distances between states across environments using pseudometrícs on S; the set of all such pseudometrics is M,, and M, is the set of metrics on probability distributions over S.
In our setting, the learner has access to a collection of training MDPs (M, drawn from p. After interacting with these environments, the learner must produce a policy π over the entire state space S, which is then evaluated on unseen MDPs from p. Similar in spirit the setting of transfer learning (Taylor& & Stone, 2009), here we evaluate the policy's 's zero-shot performance on p.
Our policy similarity metric (Section 3) builds on the concept of r-bisimulation (Castro, 2020). Under the a-bisimulation metric, the distance between two states, x and y, is defined in terms of the difference between the expected rewards obtained when following policy æ. The bibismmulation metric d satisfies a recursive equation based on the 1-Wasserstein metric Wi M + Mp, where W1 (d) B) is the minimal cost of transporting probability mass from A to B (two probability distributions on S) under the base metric d (Villani, 2008). The recursion is
dx (I, y) R (z) - Rd (v)) + 7Wi (-1x), (-
v))
K, E S.
(I)
To achieve good generalization properties, we learn an embedding function 2e : S + R* that reflects the information encoded in the policy similarity metrie; this yields policy similarity
PPeeudometrics are generalization of metrics where the distance between two distinct states can be zero.
2

---

Published as a conference paper at ICLR 2021
embedding (Section 4). We use contrastive methods (Hadsell et al., 2006; Oord et al., 2018) whose track record for representation learning is well esttbbisbd.. We adapt SimCLR (Chen et al., 2020), popular contrastive method for learning embeddings of image inputs. Given two inputs x and y. their embedding similarity is sa(x, y) = sim((z(x), ze (v)), where sim(u, v) = Tai denotes the cosine similarity function. SimCLR aims to maximize similarity between augmented versions of an image (e.g., crops, colour changes) while minimizing its similarity to other images. The loss used by
SimCLR for two versions æ, of an image, and a set X' containing other images is:
exp(Ase(1z.)) y))
Co y: x²) = - log
(2)
exp(Ase(I, v)) + En' exp(Asa(a', y))
where À is inverse temperature hyperparameter. The overall SimCLR loss is then the expected value of Ea (z, yt S), when v, y, and X' are drawn from some augmented training distribution.
3 POLICY SIMILARITY METRIC
A useful tool in learning a policy that generalizes is to understand which states result in similar behavior, and which do not. To be maximally effective, this similarity should go beyond the immediately chosen action and consider long-term behavior. In this regards, the r-bisimulation metrics are interesting as they are based on the full sequence future rewards received from different states. However, considering rewards can be both too restrictive (when the policies are the same, but the obtained rewards are not; see Figure 1.1) or too permissive (when the policies are different, but the obtained rewards not; Figure 3a). In fact, a-bisimulation metries actually lead to poor generalization in our experiments (Sections 5.1 and 5.2).
To address this issue, we instead consider the similarity between policies themselves. We replace the absolute reward difference by a probability pseudometric between policies, denoted DIST. Addi- tionally, since we would like to perform well in unseen environments, in optimal behavior. We thus use x² as the grounding policy. This yields the policy similarity metric (PSM), for which states are close when the optimal policies in these states and future states are similar. For given DIST, the PSM d= S x S + R satisfies the recursive equation
are interested in similarity
d' (xx y) = DIST(TT (TS), x- (v)) (d)) (.1z), P" (-11)
(3)
(A)
(B)
The DIST term captures the difference in local optimal behavior (A) while W1 captures long-term optimal behavior difference (B); the exact weights assigned to the are given by the discount factor. Furthermore, when DIST is bounded, d is guaranteed to be finite. While there are technically multiple PSMs (one for each DIST), we omit this distinction whenever clear from context. A proof of the uniqueness of d' is given in Proposition B.I.
Our main use of PSM will be to compare states across environments. In this context, we identify the terms in Equation 3 with specific environments for clarity and write (despite its technical inaccuracy)
d" (z. v) DIST DIST(T?((x), x5)(w) + (Wi (d') (+12), (-1w)).
PSM is applicable to both discrete and continuous action spaces. In our experiments, DIST is the total variation distance (TV) when is discrete, and we use the & distance between the mean actions of the two policies when A is continuous. PSM can be computed iteratively using dynamic programming (Ferns al., 2011) (more details in Section D. Furthermore, when x is unavailable on training environments, we replace it by an approximation A' to obtain an approximate PSM, which is close to the exact PSM depending on the suboptimality of A (Proposition D.3).
Despite resembling r-bisimulation metrics in form, the PSM possesses different characteristics which are better suited to the problem generalizing a learned policy. To illustrate this point, consider the following simple nearest-neighbot Σy := arg mingex x d' (x, y). Suppose that we use this scheme to transfer T to My, in the sense that we behave according the policy (y) = T (Z We can then bound the difference between स and Tÿ. something that is not possible d= is replaced by a TT ribisiiulatiion metric.
scheme: Given state y V. denote its closest match in X' by
Theorem 1. |Bound on policy transfer] For any y E Y, let Y ~ p* (-1
define the sequence of random states encountered starting in y y and following policy म. We have:
YG
3

---

Published as a conference paper at ICLR 2021
Byg
Σ TV (#(G),,*(3))) 120
S 1+2 d' (Z.
The proof is in the Appendix (Section B). Theorem is non-vacuous whenever d' (T., y) < 1/(1+>)) In particular, d' (ix 3) = 0 implies that the transferred policy is optimal. Although this scheme is not practical (computing d' requires knowledge of Tÿ)) shows that meaningful policy generalization can be obtained if we can find a mapping that generalizes across S. Put another way. PSM gives us principled way of lifting generalization across inputs (a supervised learning problem) to generalization across environments. We now describe how to employ PSM for learning representations that put together states in which the agent's long-term optimal behavior similar.
4 LEARNING CONTRASTIVE METRIC EMBEDDINGS
To generalize a learned policy to new envi- ronments, we build on the success of con- trastive representations (Section 2). Given a state similarity metric we develop a general procedure (Algorithm 1) to learn contrastive metric embeddings (CMEs) for d. We utilize the metric d for defining the set of positive and negative pairs, as well assigning importance weights to these pairs in the contrastive loss (Equation 4). 5: end for
Algorithm 1 Contrastive Metric Embeddings (CMEs)
1: Given: State embedding zą (-), Metric d(- .) Train-
ing environments {M:); } Hyperparameters: Temperature 1/X, Scale 3, Total training steps K
2: for each step k 1..... K do
3: Sample pair of training MDPs Mr,
4: Update to minimize LCME where
LCME EM (L My]]
We first apply transformation convert d to similarity measure Т, bounded in [0, 1] for "soft" similarities. In this work, we transform using the Gaussian kernel with positive scale parameter J3, B controls the sensitivity of the similarity measure to d.
that is, [2x y)
(ppp-d(x,,)))
Second, we select the positive and negative pairs given a set of states X' S X and y from MDPs Mx, respectively. For each anchor state y V, we use its nearest neighbor in X' based on the similarity measure I' to define the positive pairs {(Tg. v)), where Σy = argmax T(x, y). The #EX'
remaining states in 2', paired with y, are used as negative pairs. This choice of pairs motivated by Theorem which shows that if we transfer the optimal policy in M, the nearest neighbors defined using PSM, its performance in My has suboptimality bounded by PSM.
Next, we define a soft version of the SimCLR contrastive loss (Equation 2) for learning the function 28, which maps states (usually high- digh-dimmnsional) to embeddings. Given a positive state pair (H y). the set X", and the similarity measure I, the loss (pseudocode provided in Section J.2) given by
I I( v) exp(Ase v)) [(x', y)) exp (ssp(ss(x''))) v))
Co ys x') log T( y) exp(Aso (Ty, v)) Ex'x
(4),
where we the same notation in Equation Following SimCLR, we use non-linear projection of the representation as zo (Figure A. The agent's policy is an affine function of the representation.
The total contrastive loss for Mx and My utilizes the optimal trajectories x { and 3 = {y I. where 2441 P (- 22)) and V1t1 PJ (.1ve). We set X" x and define
Lo(Mx.M5) = Ey-rr 16o xx) where i. = argmax [( (z,
#Er
We refer to CMES learned with policy similarity metric as policy similarity embeddings (PSEs). PSEs can trivially be combined with data augmentation by using augmented states when computing state embeddings. We simultaneously learn PSEs with the RL agent by adding LCmE (Algorithm 1) as an auxiliary objective during training. Next, we illustrate the benefits from this auxiliary objective.
JUMPING TASK FROM PIXELS: A CASE STUDY
Task Description. The jumping task (Tachet des Combes et al., 2018) (Figure 1) captures, using well-defined factors of variations, whether agents can leam the correct invariances required for generalization, directly from image inputs. The task consists of an agent trying to jump over an obstacle. The agent has access to two actions: right and jump. The agent needs to time the jump
4

---

Published as a conference paper at ICLR 2021
Table 1: Percentage (%) of test tasks solved by different methods without and with data augmentation. The "wide", "narrow", and random grids are described in Figure We report average performance across 100 runs with different random initializations, with standard deviation between parentheses.
Data Augmentation
Grid Configuration (%)
Method
"Wide"
"Narrow"
Random
Dropout and Eg reg.
17.8 (2.2)
10.2 (4.6)
9.3 (5.4)
X
Bisimulation Transfer* PSEs
(0.0)
17.9 (0.0)
30.9 (4.2)
33.6 (10.0)
9.3 (5.3)
37.7 (10.4)
RandConv
50.7 (24.2)
33.7 (11.8
71.3 (15.6)
RandConv + Bisimulation
(17.6)
17.4 (6.7)
(15.6)
RandConv + PSEs
87.0 (10.1)
52.4 (5.8)
83.4 (10.1)
Obstacle Position
i
Rloor
Height
(a) "Wide" grid
(b) "Narrow" grid
(c) Random grid
Figure 2: Jumping Task: Visualization of average performance of PSEs with data augmentation across different configurations. We plot the median performance across 100 runs. Each tile in the grid represents a different task (obstacle position/floor height combination). For each grid configuration, the height varies along the y-axis (1) heights) while the obstacle position varies along the z-axis (26 locations). The red letter T indicates the training tasks. Random grid depicts only one instance, each run consisted of different test/train split. Beige tiles are tasks PSEs solved while black tiles are tasks PSEs did not solve when used with data augmentation. These results were chosen across all the 100 runs to demonstrate what the average reported performance looks like.
precisely, at specific distance from the obstacle, otherwise it will eventually hit the obstacle. Different tasks consist in shifting the floor height and/or the obstacle position. To generalize, the agent needs to be invariant to the floor height while jump based the obstacle position. The obstacle can be in 26 different locations while the floor has different heights, totaling 286 tasks.
Problem Setup. We split the problem into 18 seen (training) and 268 unseen (test) tasks to stress test generalization using a few changes in the underlying factors of variations seen during training. The small number of positive examples results in highly unbalanced classification problem with low amounts data, making it challenging without additional inductive biases. Thus, we evaluate gener- alization in regimes with and without data augmentation. The different grids configurations (Figure 2) capture different types generalization: the "wide" grid tests generalization via "interpolation", the "narrow" grid tests -attof-distribution generalization via "extrapolation" and the random grid instances evaluate generalization similar to supervised learning where train and test samples are drawn i.i.d. from the same distribution.
We used RandConv (Lee et al., 2020a), a oftc-oofthhearr data augmentation for generalization. For hyperparameter selection, we evaluate all agents on a validation set containing 54 unseen tasks in the "wide" grid (Figure 2a) and pick the parameters with the best validation performance. We use these fixed parameters for all grid configurations to show the robustness of PSEs hyperparameter tuning. We first compute the optimal trajectories in the training tasks. Using these trajectories, we compute PSM using dynamic programming (Section D.I). We train the agent by imitation learning, combined with an auxiliary loss for PSEs (Section 4). More details are in Section G.
5.1
EVALUATING GENERALIZATION ON JUMPING TASK
We show the efficacy of PSEs compared to common generalization approaches such regulariza- tion (e.g., Cobbe et al., 2019; Farebrother et al., 2018), 8), and data augmentation (e.g., Lee et al., 2020a; Laskin et al., 2020a), which are quite effective on pixel-based RL tasks. We also contrast PSEs with bisimulation transfer (Castro& & Precup, 2010), tabular state-based transfer approach based on bisim- ulation metrics which does not do any learning and bisimulation preserving representations (Zhang et al., 2021), 1), showing the advantages of PSM over prevalent state similarity metric.
We have 18 different trajectories with several examples for the action right, but only one instance of jump action per trajectory, leading just 18 total instances of the action jump.
5

---

Published as a conference paper at ICLR 2021
(c) PSM & embed- dings
(a) Optimal Trajectories
(b) PSES
(d) ---hiimm. CMEs
Figure 3: Embedding visualization. (a) Optimal trajectories on original jumping task (visualized as coloured blocks) with different obstacle positions. We visualize the hidden representations using UMAP, where the color of points indicate the tasks of the corresponding observations. Points with the same number label correspond to same distance of the agent from the obstacle, the underlying optimal invariant feature across tasks.
We first investigated how well PSEs generalize over existing methods without incorporating additional domain knowledge during training. Table summarizes, the setting without data augmentation, the performance of these methods in different train/test splits (c.f. Figure 2 for detailed description). PSEs, with only 18 examples, already leads to better performance than standard regularization.
PSEs also outperform bisimulation transfer in the "wide" and random grids. Although bisimulation transfer is impractical when evaluating zero-shot generalization, we still performed this comparison, unfair to PSEs, to highlight their efficacy. PSEs perform better because, in contrast to bisimulation, PSM is reward agnostic (c.f. Proposition D.1) the expected return of the jump action is quite different depending on the obstacle position (c.f. Figure G.2 for visual juxtaposition of PSM and bisimulation). Overall, these results are promising because they place PSEs as an effective generalization method that does not rely on data augmentation.
Nevertheless, PSEs are complementary to data augmentation, which consistently improves gener- alization in deep RL. We compared RandConv combined with PSEs to simply using RandConv. Domain-specific augmentation also succeeds in the jumping task. Thus, it is not surprising that RandConv is so effective compared to techniques without augmentation. Table (2nd row) shows that PSEs substantially improve the performance of RandConv across all grid configurations. More- over, Table (2nd row) illustrates that when combined with RandConv, bisimulation preserving representations (Zhang et al., 2021) diminish generalization by 30 50% relative to PSEs.
Notably, Table (1*t row) indicates that learning- based methods are ineffective on the "narrow" grid without data augmentation. That said, PSEs do work quite well when combined with RandConv. However, even with data augmentation, generalization vicinity of training tasks, exhibiting the challenge this grid poses for learning based methods. We believe this is due to the poor extrapolation ability of neural networks (e.g.. Haley & Soloway, 1992; Xu et al., 2020), which is more perceptible without prior inductive bias from data augmentation.
"narrow" grid happens only around the
5.2
UNDERSTANDING GAINS FROM PSES: ABLATIONS AND VISUALIZATIONS
PSEs are contrastive metric embed- dings (CMEs) learned with PSM. We investigate the gains from CMES (Section 4) and PSM (Sec- tion 3) by ablating them. CMEs can be learned with any state similarity metric bisimulation (Castro, 2020) as an al- ternative. Similarly, PSM be used with any metric embedding we use egeemeddisggs (Section E) as an alternative, which Zhang et al.
Table 2: Ablating PSEs. Percentage (%) of test tasks solved when
we ablate the similarity metric and lcarning procedure for metric embeddings in the data augmentation setting wide grid. PSEs, which combine CMEs with PSM, considcrably outperform other embeddings. We report the average performance across runs with standard deviation between parentheses. All ablations except PSEs deteriorate performance compared to just using data augmentation (RandConv), reported Table 1.
Metric Embedding
E2-embeddings
CMEs
π bisimuulatio
41.4 (17.6)
23.1 (7.6)
PSM
17.5 (8.4)
87.0 (10.1)
Bisimulation transfer assumes oracle access to dynamics and rewards unseen environments as well as tabular state space compute the exact bisimulation metric (Section C).

---

Published as a conference paper at ICLR 2021
(2021) employed with r'-bisimulation for learning representations in single task RL setting. For fair comparison, we tune hyperparameters for 128 trials for each ablation entry in Table 2.
Table 2 show that PSEs (= PSM CMEs) generalize significantly better than n -isimuulation with CMEs or fz-embeddings,, both of which significantly degrade performance (-60% and -45%, respectively). This is expected, as x r""bisimulation imposes the incorrect invariances for the jumping task (c.f. Figures G.2a and G.2d). Additionally, looking at the rows of Table 2, CMEs are superior than fg-embeddings for PSM while inferior for x r'-bisimulation. This result is in line with the hypothesis that CMEs better enforce the invariances encoded by similarity metric compared to fg-embeddings (c.f. Figures 3b and 3c).
Visualizing learned representations. We visualize the metric embeddings in the ablation above by projecting them two dimensions with UMAP (MeInnes al., 2018), a popular visualization technique for high dimensional data which better preserves the data's global structure compared to other methods such as t-SNE (Coenen Pearce, 2019).
Figure 3 shows that PSEs partition the states into two sets: (1) states where single suboptimal action leads to failure (all states before jump) and (2) states where actions do not affect the final outcome (states after jump). Additionally, PSEs align the labeled states in the first set, which have PSM distance of zero. These aligned states have the same distance from the obstacle, the invariant feature that generalizes across tasks. On the other hand, fz-embeddings (Zhang et al., 2021) with PSM do not align states with zero PSM except the state with the jump action poor generalization, observed empirically, is likely due states with the same optimal behavior ending up with distant embeddings. CMEs with x -bisimulation align states with π -bisimullttion distance zero such states are equidistant from the start state and have different optimal behavior for any pair of tasks with different obstacle positions (Figure G.2c).
5.3 JUMPING TASK WITH COLORS: WHERE TASK-DEPENDENT INVARIANCE MATTERS
PSEs capture invariances that are usually orthogonal to task-agnostic invariances from data augmentation. This difference is important because, for certain tasks, data augmentation erroneously alias states with different optimal behavior. Domain knowledge is often required to select appropriate augmentations, otherwise augmen- tations can hurt generalization. In contrast, PSEs do require any domain knowledge but instead exploit the inherent structure of the RL tasks.
Method
RandConv
Dropout
and l2 reg.
RandConv
PSEs
PSEs
10
20
30
40
To demonstrate the difference between PSEs and data augmentation, we simply include colored obstacles in the jumping task (see Figure G.5). In this modified task, the optimal behavior of the agent depends on the obstacle color: the agent needs to jump over the red obstacle but strike the green obstacle to get a high return. The red obstacle task has the same difficulty as the original jumping task while the green obstacle task is easier. We jointly train the agent with 18 training tasks each, for both obstacle colors, and evaluate generalization on unseen red tasks.
% Test tasks solved
Figure 4: Percentage (%) of red obstacle test tasks solved when trained, jointly with red and green obstacles, on the "wide" grid. We report the mean across 100 runs. Error hars show 99% confidence interval for the mean.
Figure 4 shows the large performance gap between PSEs and data augmentation with RandConv. All methods solve the green obstacle tasks (Table G.1). As opposed to the original jumping task (c.f. Ta- ble 1), data augmentation inhibits generalization since RandCony forces the agent to ignore color, conflating the red and green tasks (Figure G.6). PSEs still outperform regularization and data aug- mentation. Furthermore, data augmentation performs better when combined with PSEs. Thus, PSEs are effective even when data augmentation hurts performance.
5.4 EFFECT OF POLICY SUBOPTIMALITY ON PSES
To understand the sensitivity of learning effective PSEs to the quality the policies, we compute PSEs using e-suboptimal policies on the jumping task, which take the optimal action with probability 1- c and the subopotimal action with probability e. We evaluate the generalization performance of
PSEs for increasingly suboptimal policies. ranging from the optimal policy (e = 0) to the uniform random policy (e = 0.5). To isolate the effect of suboptimality on PSEs, the agent still imitates the optimal actions during training for all c.
7

---

Published as a conference paper at ICLR 2021
Figure 5 shows that PSEs show near-optimal generalization with é I 0.4 while degrade generalization with an uniform random pol- icy. This result is well-aligned with Proposition D.3, which shows that for policies with decreasing suboptimality, PSM approxima- tion becomes more accurate, resulting in improved PSEs. Overall, this study confirms that the utility of PSEs for generalization is robust to suboptimality. One reason for this robustness is that PSEs are likely to align states with similar long-ten tingttrmm greedy op- timal actions, resulting in good performance even with suboptimal policies that preserve these greedy actions.
% Test tasks solved
100
75
50
25
0.00.10.20.3..40.5
Suboptimality (c)
Figure 5: % of test tasks solved using PSEs computed using e- suboptimal policies on the "wide" grid. We report the mean across 100 runs. Error hars show one standard deviation.
ADDITIONAL EMPIRICAL EVALUATION
In this section, we exhibit that PSM ignores spurious information for generalization using LQR task (Song et al., 2019) with non-image inputs. Then, demonstrate the scalability of PSEs without explicit access to optimal policies in an RL setting, with continuous actions, using Distracting DM Control Suite (Stone al., 2021; Zhang al., 2018b).
6.1 LQR WITH SPURIOUS CORRELATIONS
We show how representations learned using PSM, when faced with semantically equivalent environ- ments, can learn the main factors of variation and ignore spurious correlations that hinder generaliza-
LQR Cost relative to Oracle
30
20
Num Distractors
Th 500
tion. We use LQR with distractors (Song et al., 2019;
10
Sonar et al., 2020) to assess generalization in a feature- based RL setting with linear function approximation. The distractors are input features that are spuriously correlated with optimal actions and can be used for predicting these actions during training, but hurt gen- cralization. The agent learns a linear policy using 2 environments with fixed distractors. This policy is evaluated on environments with unseen distractors.
Yig 1000
The 10000
0.01
Over- IPO Sparsity PSM parameterization Method
Figure LQQ generalization: Absolute test er- ror LQR cost relative to the oracle (which has access to true state), of various methods trained with n& distractors on 2 training environments. Lower error is better. PSM achieves near-optimal performance. We report mean across 100 differ- ent seeds. Error bars (non-cxistent for most meth- ods) show 99% confidence interval for mcan.
We aggregate state pairs in training environments with near-zero PSM. We contrast this approach with (i) IPO (Sonar al., 2020), a policy optimization method based on IRM (Arjovsky et al., 2019) for out- of-distribution generalization, (ii) overparametrriztion, which leads to better generalization via implicit regularization (Song et al., 2019), and (iii) weight sparsity using é -eeegularization since the policy weights which generalize in this task are sparse.
All methods optimally solve the training environments; however, the baselines perform abysmally in terms of generalization compared state aggregation with PSM (Figure 6), indicating their reliance on distractors. PSM obtains near-optimal generalization which we corroborate through this conjecture (Section H.1): Assuming zero state aggregation error with PSM, the policy learned using gradient descent is independent of the distractors. Refer to Section H for a detailed discussion.
6.2 DISTRACTING DM CONTROL SCITE
Finally, we demonstrate scalability of PSES on the Distract- ing DM Control Suite (DCS) (Stone et al., 2021), which tests whether agents can ignore high-dimensional visual distractors irrelevant to the RL task. Since we do not have access to opti- mal training policies, we use learned policies as proxy for T for computing PSM as well as collecting data for optimizing PSEs. Even with this approximation, PSEs outperform state- of-the-art data augmentation. DCS extends DM Control (Tassa et al., 2020) with visual distractions. We use the dynamic back- ground distractions (Stone et al., 2021; Zhang et al., 2018b) where video is played in the background from specific frame. The video and the frame are randomly sampled every new episode. We use 2 videos during training (Figure 7) and evaluate generalization 30 unseen videos (Figure I.1).
Figure 7: Distracting Control Suite: Snapshots of training environments.
8

---

Published as a conference paper at ICLR 2021
Table 3: Generalization performance with unseen distractions in the Distracting Suite at 500K steps. We report the average scores across 5 seeds standard error. All methods are added SAC (Haarnoja et al., 2018).
Pretrained initialization uses DrQ trained for 500K steps. Figures 1.2 and 1.3 show generalization curves.
Initialization
Method
BiC-catch
C-swingup
C-run
F-spin
R-easy
W-walk
DrQ DrQ + PSEs
747+288
582+-42
220+-12
646++54
931+114
549++83
Random
821=117
749:119
308:12
779::49
955110
789-228
DrQ
748+:30
6894-22
219+110
764-1488
943+-10
709-:29
Pretrained
DrQ + PSEs
805:225
733++11
282:1-8
803--199
962-1:11
829--21
All agents are built top of SAC (Haarnoja al., 2018) combined with DrQ (Kostrikov et al., 2020), an augmentation method with state-of-the thrtth-art performance on DM control. Notably, DrQ outperforms CURL (Laskin et al., 2020b), a strong contrastive method for RL. Without data augmentation on DM control, SAC performs poorly, even during training (Kostrikov et al., 2020). We augment DrQ with an auxiliary loss for learning PSEs and compare it with DrQ (Table 3). Orthogonal to DrQ, PSEs align representations of different states across environments based on PSM (c.f. Figure A. All agents are trained for 500K environment steps with random crop augmentation. For computing PSM, we use policies learned by DrQ pretrained on training environments for 500K steps.
First, we investigate how much better PSEs generalize relative to DrQ, assuming the agent provided with PSM beforehand. The agent's policy is randomly initialized so that additional gains over DrQ can be attributed to the auxiliary information from PSM. The substantial gains in Table indicate that PSEs are more effective than DrQ for encoding invariance to distractors.
Since PSEs utilize PSM approximated using pretrained policies, we also compare to a DrQ agent where we initialize it using these pretrained policies. This comparison provides the same auxiliary information to DrQ as available to PSEs, thus, the generalization difference stems from how they utilize this information. Table 3 demonstrates that PSEs outperform DrQ with pretrained initialization, indicating that the additional pretraining steps are more judiciously utilized for computing PSM as opposed to just longer training with DrQ. More details, including learning curves, are in Section I.
RELATED WORK
PSM (Section 3) is inspired by bisimulation metrics (Section C). However, different than traditional bisimulation (e.g.. Larsen Skou, 1991; Givan et al., 2003; Ferns et al., 201 PSM more tractable as it defined with respect to a single policy similar to the recently proposed x -iiiimiulattion (Cas- tro, 2020; Zhang et al., 2021). However, in contrast to PSM, bisimulation metrics rely on reward information and may not provide a meaningful notion of behavioral similarity in certain environ- ments (Section 5). For example, states similar under PSM would have similar optimal policies, yet can have arbitrarily large T bbisimulltiinn distance between them (Proposition D.1).
PSEs (Section 4) use contrastive learning to encode behavior similarity (Section 3) across MDPs. Previously, contrastive learning has been applied for imposing state self-consistency (Laskin et al., 2020b), capturing predictive information (Oord et al., 2018; Mazoure et al., 2020; Lee et al., 2020b) or encoding transition dynamics (van der Pol et al., 2020; Stooke al., 2020; Schwarzer al., 2020) within an MDP. These methods can be integrated with PSEs to encode additional invariances. Interestingly, in a similar spirit PSEs, Pacchiano et al. (2020); Moskovitz et al. (2021) explore comparing behavioral similarity between policies to guide policy optimization within an MDP.
PSEs are complementary data augmentation methods (Kostrikov et al., 2020: Lee et al., 2020a; Raileanu et al., 2020; Ye et al., 2020), which have recently been shown to significantly improve agents' generalization capabilities. In fact, we combine PSEs to state-of-the-arr augmentation methods including random convolutions (Lee et al., 2020a; Laskin et al., 2020a) in the jumping task and DrQ (Kostrikov et al., 2020) on Distracting Control Suite, leading to performance improvement.
CONCLUSION
This paper advances generalization in RL by two contributions; (1) the policy similarity metric (PSM) which provides a new notion of state similarity based on behavior proximity, and (2) contrastive metric embeddings, which harness the benefits of contrastive learning for representations based similarity metric. PSEs combine these two ideas to improve generalization. Overall, this paper shows the benefits of exploiting the inherent structure in RL. for learning effective representations.

---

Published as a conference paper at ICLR 2021
REFERENCES
Rishabh Agarwal, Chen Liang, Dale Schuurmans, and Mohammad Norouzi. Learning generalize from sparse and underspocified rewards. In /CML, 2019.
Martín Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Pas Paz. Invariant risk minimization. CoRR. abs/1907.02835 2019.
Sanjeev Arora, Nadav Cohen. Wei Hu, and Yuping Luo. Implicit regularization deep matrix factorization, In Advances Neural Information Processing Systems, pp. 7413-7424, 2019.
Yogesh Balaji, Swami Sankaranarayanan, and Rama Chellappa. Metareg: Towards domain generalization using metta-rggularization. In Advances in Neural Information Processing Systems, pp. 998-1008, 2018.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic Markov decision processes, In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2020.
Pablo Samuel Castro and Doina Precup. Using bisimulation for policy transfer in MDPs. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2010.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. CoRR, abs/2002.05709, 2020.
Karl Cobbe, Oleg Klimov. Christopher Hesse. Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. In Proceedings of the International Conference on Machine Learning (ICML), 2019.
Andy Coenen and Adam Pearce. Understanding umap, 2019. URL https:// http:pp::////aiiieee code. tttp:paiiiiie..iii.........aaab. in/ understane -estaandnnnnnaaann unap.
Jesse Farebrother, Marlos C. Machado, and Michael Bowling, Generalization and regularization in DQN. In
NewrlPS Deep Reinforcement Learning Workshop, 2018
Norm Ferns Prakash Panangaden, and Doina Precup. Metries for finite Markov decision processes. In Proceedings of the Conference in Uncertainty Artificial Intelligence (UAI), 2004.
Norm Ferns, Pablo Samuel Castro, Doina Precup, and Prakash Panangaden. Methods for computing state similarity Markov decision processes. Proceedings of the 22nd Conference on Uncertainty in Artificial Intelligence, UAI '06. AUAI Press, 2006.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes, SIAM Journal on Computing, 4066): 2011.
Norman Ferns and Doina Precup. Bisimulation metrics are optimal value functions. In The 30th Conference on Uncertainty in Artificial Intelligence, pp. 10, 2014.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep
networks. In Proceedings of the International Conference on Machine Learning (ICML), 2017.
Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization in markov decision processes. Artificial Intelligence, 147(1-2):163-22)) 2003.
Daniel Golovin, Benjamin Solnik, Subhodeep Moitra, Greg Kochanski, John Karro, and D Sculley. Google vizier: A service for black-box optimization. In Proceedings of the 23rd ACM SIGKDD international conference knowledge discovery and data mining, pp. 1487 1495, 2017.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, pp. 6151-6159. 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critie: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. arXiv preprint arXiv:1801.01290 2018.
Raia Hadsell, Sumit Chopra, and Yann LeCun. Dimensionality reduction by learning an invariant mapping. In 2006 IEEE Computer Society Conference on Computer Vision and Puttern Recognition (CVPR'06), volume 2, pp. 7735-1742. IEBE, 2006.
Pamela J Haley and DONALD Soloway. Extrapolation limitations of multilayer feedforward neural networks. In (Proceedings 1992] IJCNN International Joint Conference Neural Networks, volume 4, pp. 25-30. IERR, 1992.
10

---

Published as a conference paper at ICLR 2021
Maximilian Igl. Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam Devlin, and Katja Hofmann. Generalization in reinforcement learning with selective noise injection and information bottleneck. In Advances in Neural Information Processing Systems, pp. 13978-13990, 13990, 2019.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Domain extrapolation via regret minimization, arXiv preprint arXiv:2006. 03908, 2020.
Arthur Juliani, Ahmed Khalifa, Vincent-Pierne Berges, Jonathan Harper, Ervin Teng, Hunter Henry, Adam Crespi, Julian Togelius, and Danny Lange. Obstacle Tower: A generalization challenge in vision, control, and planning. In Proceedings of the International Joint Conference on Artificial Intelligence (IJCAI), 2019.
Niels Justesen, Ruben Rodriguez Torrado, Philip Bontrager, Ahmed Khalifa, Julian Togelius, and Sehastian Risi. Illuminating generalization deep reinforcement learning through procedural level generation. arXiv preprint arXiv: 10729, 2018.
Taylor Killian, George Dimitri Konidaris, and Finale Delei-Velez. Robust and efficient transfer learning with hidden parameter Markov decision processes. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), pp. 4949-499999 4950, 2017.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforeement learning from pixels. CoRR, abs/2004. 100004113444 13649, 2020.
Kim G Larsen and Arne Skou. Bisimulation through prohabilistic testing. Information and computation, 94(1): 1-28, 1991.
Michael Laskin, Kimin Lec, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. CoRR, abs/2004. 14990, 2020a.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. Proceedings of the 37th International Conference on Machine Learning, Vienna, Austria, PMLR 119, 2020b. arXiv:2003.06417.
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lec. Network randomization: A simple technique for general- ization in deep reinforcement learning. In The International Conference on Learning Representations (ICLR), 2020a.
Kuang-Huei Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama. Predictive information accelerates learning in rl. arXiv preprint arXiv:2007.12401, 2020b.
Da Li, Yongxin Yang. Yi-Zhe Song, and Timothy M. Hospedales. Learning to generalize: Meta-lcarning for domain generalization. Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2018.
Bogdan Mazoure, Remi Tachet des Combes, Thang Long DOAN, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. Advances in Neural Information Processing Systems, 33, 2020.
McInnes, J. Healy, and Melville. UMAP: Uniform Manifold Approximation and Projection Dimension Reduction, ArXiv e-prints, February 2018.
Ted Moskovitz, Michael Arbel, Ferenc Huszar, and Arthur Gretton. Efficient wasserstein natural gradients for reinforcement learning. In International Conference on Learning Representations, 2021.
Junhyuk Oh, Satinder P. Singh, Honglak Lee, and Pushmeet Kohli. Zero-shot task generalization with multi-task deep reinforcement learning. In Proceedings of the International Conference on Machine Learning (ICML), 2017.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 1rXtvv/807.03748, 2018
Aldo Pacchiano, Jack Parker-Holder, Yunhao Tang, Krzysztof Choromanski, Anna Choromanska, and Michael Jordan, Learning to score hehaviors for guided policy optimization. International Conference on Machine Learning, 2020.
Charles Packer, Katelyn Gao, Jemnj Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning. arXiv preprint 1800 18t..2222.....2... 2018
Christian F. Perez, Felipe Petroski Such, and Theofanis Karaletsos. Generalized hidden parameter mdps transferable model-hased rl in handful of trials. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2020.
11

---

Published as a conference paper at ICLR 2021
Jordi Pont-Tuset, Federico Perazzi. Sergi Caelles, Pablo Arbeláez, Alex Sorkine-Hoomung. and Luc Van Gool. The 2017 davis challenge on video object segmentation. arXiv preprint arXiv: 17X4vv1704.00000 00675, 2017.
Martin L Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. John Wiley & Sons, Inc., 1994.
Roberta Raileanu, Max Goldstein, Denis Yarats, llya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in deep reinforcement learning. arXiv preprint arXiv:2006. /2862, 2020.
Aravind Rajeswaran, Kendall Lowrey, Emanuel Todorov, and Sham M. Kakade. Towards generalization and simplicity in continuous control. Advances Neural Information Processing Systems (NeurIPS), 2017.
Benjamin Recht, tour reinforeement learning: The view from continuous control. Annual Review of Control, Robotics, and Autonomous Syxtems, 2019.
Max Schwarzer, Ankesh Anand, Rishab Goel, Devon Hjelm, Aaron Courville, and Philip Bachman. Data- efficient reinforcement learning with momentum predictive representations. arXiv preprint arXiv:2007.05929, 2020.
Anoopkumar Sonar, Vincent Pacelli, and Anirudha Majumdar. Invariant policy optimization: Towards stronger generalization in reinforcement learning. arXiv preprint arXiv:2006.010 01096, 2020.
Xingyou Song, Yiding Jiang, Yilun Du, and Behnam Neyshabur. Observational overfitting in reinforcement learning. In The International Conference Learning Representations (ICLR), 2019.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite a challenging benchmark for reinforcement learning from pixels. arXiv preprint urXiv:2101.02722 2021.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. arXiv preprint 2:Xivv2009.083111 2020.
Remi Tachet des Combes, Philip Bachman, and Harm van Seijen. Learning invariances for policy generalization. In Workshop track at the International Conference on Learning Representations (ICLR), 2018.
Yujin Tang, Duong Nguyen, and David Ha. Neuroevolution of self eetf-rnterpeetable agents. arXiv preprint arXiv:2003.08 081660000066........ 2020.
Yuval Tassa, Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap. and Nicolas Heess. iin_control: Software and tasks for continuous control. arXiv preprint arXiv:2006. 12983, 2020.
Matthew E. Taylor and Peter Stone. Transfer learning for reinforcement learning domains: A survey. Journal of Machine Learning Research, 10:1633-1685. 2009.
Josh Tobin, Rachel Fong. Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomiza- tion transferring deep neural networks from simulation to the real world. In 2017 JEEE/RSJ International Conference on Intellgents Robots and Systems (IROS), pp. 23-30. IEEF, 2017.
Elisc van der Pol, Thomas Kipf, Frans Olichock, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. In Proceedings of the 19th International Conference on Autonomous Agents and MultiAgent Systems, pp. 1431-1439, 2020.
Cédric Villani. Optimal transport: old and new. Springer, 2008
Sam Witty, Jun Ki Lec, Emma Tosch, Akanksha Atrey, Michael Littman, and David Jensen. Measuring and characterizing generalization in deep reinforcement learning. In NewrIPS Critiquing and Correcting Trends in Machine Learning Workshop, 2018
Keyulu Xu, Jingling Li, Mozhi Zhang, Simon S Du Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward graph neural networks. arXiv preprint ar arXvv:2009. 1148, 2020.
Chang Ye, Ahmed Khalifa, Philip Bontrager, and Julian Togelius. Rotation, translation, and cropping for zero-shot generalization. CoRR, Abs/2001.09990 2020.
Amy Zhang, Nicolas Ballas, and Joelle Pineau. A Dissection of Overfitting and Generalization in Continuous Reinforcement Leaminng CoRR, abs/1806.000 1806007937 2018a.
Amy Zhang, Yuxin Wu, and Joelle Pineau. Natural environment benchmarks for reinforcement learning. arXiv preprint arXiv 18/ 1.06032, 20188 8b.
12

---

Published as a conference paper at ICLR 2021
Amy Zhang, Clare Lyle. Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In Proceedings of International Conference on Machine Learning (ICML), 2020.
Amy Zhang, Rowan McAllister, Roberto Calandra. Yarin Gal, and Sergey Levine. Learning invariant repre- sentations for reinforcement learning without reconstruction. The International Conference on Learning Representations (ICLR), 2021.
Chiyuan Zhang, Oriol Vinyals, Rémi Munos, and Samy Bengio. study on overfitting in deep reinforvement learning. CoRR, abs/ 804,0683, 2018c.
13

---

Published as a conference paper at ICLR 2021
Appendix
A
LEARNING CONTRASTIVE METRIC EMBEDDINGS
Augmentation
Representation
Projection
Input
f
hg
f
Z
X
Contrastive Learning
y
f
h
y
W
Reinforcement Learning
π
Policy
Figure A.I: Architecture for learning CMEs. Given input pair (st, y), first apply the (optional) data augmentation operator produce the input augmentations V == 10 (x); V := (y)) When not using data augmentation. V is equal to the identity operator, that is, Væ (x) = z. The agent's policy network then outputs the representations for these augmentations by applying the encoder fe. that is, fa = fM( fy (,) These representations (x) ho(fz) (f=), 20 ho(f)) These metric embeddings are trained using the contrastive loss defined in Equation (4). The policy Te is an affine function of the representation, that is, Ta(-ly) W W' f,
projected using non-linear projector he obtain the embedding ze. that
b, where W, b are learned weights and biases. The entire network is trained end-to-end jointly using the reinforcement auxiliary contrastive loss.
learning (or imitation learning) loss conjunction
B PROOFS
We begin by defining some notation which will used throughout these results:
We denote Ez2ol7 T (T)) = Ey; EEeo (VV(())))))) (x))))
For any E Y, let Y ~ Pa (-18,-1) where y y.
TV"(Y") Bosten
(V
We now proceed with some technical lemmas necessary for the main result.
Lemma 1. Given any two pseudometrics' d, d' M and probability distributions Px, Py where
X, X, C S, we have:
(d)( Px) [d d'll (d)
Proof. Note that the dual of the linear program for computing W (d) Py) is given by
Σ T(x, y) v)
min
ZEN, DEY
subject to (((,,))= = Po(v) Vy. ΣT(,,)) = Px(x) Vz. T(r,y) 2 Vv,
Psecddometrics are generalization of metrics where the distance between distinct states be
14

---

Published as a conference paper at ICLR 2021
Using the dual formulation subject to the constraints above, W (d) can be written
w (1)(xx Py) l|d d'I (d d' d') Py) S ||d - d'|l
-mim Σ
y( (d(z, v) - d' dz v) + d (z.
T a vEy
S min Σ
T( (lld d' d' d'(r, y))
aEA; yEY
= ||d d'I + min
Σ
T(x, v) d' (2,,g)
ZEX, BEY
|| d' W (dx)
Lemma 2. Given any yo E Y, we have:
2 p (y1130) p= (91180)) (x))) $ (V((())), T'()))
Σ
1-7
VIY
Proof.
Σ
ps (3)1100) - -p" (viluo)) (V)) s
Σ
pa (")1100)) - p= (w1110) TV
MEEY
BLEY
Ply a) (#(aly0) - 7 (Tlo0)) TV
EIEY
P(y 1yo a) [[alvo)) - (alyo))
VIEY aEA
Σ (aab0)) - n (aluo)l
Σ
P(y a)
aEA
MEY
[T(ayo)) - T (aluo))
1 (V . ((0o)
Lemma 3. Given any yo E Y, (TT"
< 1 dt (Tn v1), we have:
1 + 7 Wi (d') (pr' (-17yo); pm (-180)
Σ pd (v
S
VIEY
Proof. Note that we have the following equality, where 0 is a vector zeros:
Σ p Pr (x)
Σ p" (V) Σ p (o))
VIEY
VIEY
BEX
which is the same form as the primal LP for W Widd") (-110), pr PT (-) By assumption, we have that
(V)
+ 7 < 1-
d' (ăg,,
This implies that 1-7 TV (Y:") a feasible solution to Wi (id' (-180), p (-1ago)): 1+7
Σ pm (v1100)
1+ 1- TV" S (dd")) pm" (-17 p= p= (-1w0)
VIEY
and the result follows.
15

---

Published as a conference paper at ICLR 2021
Proposition B.1. The operator FF given by:
F(d)(a,)) = DIST(" (z), (V)) + ( (1x), PS (-lu))
a contraction mapping and has a unique fixed point for a bounded dist.
Proof. We first prove that F is contraction mapping. Then, simple application of the Banach Fixed Point Theorem asserts that F has a unique fixed point. Note that for all pseudometrics d, d' E M,, and for all states E X, EV.
F(d)(a,)) - (z.y) =7 (M) (11z); pe (-1)) - (idd') (1p), P3 (lw))
Lemma 1 S
7 (lld d'l + (idd pj (-19) (dd) (-b))
=7lld-d'l
Thus, F(dd)) S 7 d'||, so that F is contracting mapping for y 1 and has unique fixed point d'
Theorem 1. [Bound on policy transfer] For any y E V, let y p* (* define the sequence of random states encountered starting in Y and following policy ă. We have:
Byz
E (f(Y)), x (V,)) 120
1 (Tg.
Proof. We will prove this by induction. Assuming that the bound holds for TV", we prove the bound holds for TV=+1, The base case for = follows from (V), a (v)) = TV (ar (*",,, x (v)) S
d' (Z 3). Note that TV < -_ since the TV distance per time-step can be at most 1.
Let P. (y/ly) denote the probability of ending in state y' E Y after steps when following policy and starting from state y. We then have:
TV=+l(y)))
Σ KY
( x +vx))+
Σ PK (skl1)
TV( + E
p (ye+ (
BEEY
DEFIEY
Σ BYEY
P (uklM)
(V(#(v),, x0 (ve)) +
( (BK4t1B4)) - p= (yx (wx=) (x.
DE+IEY
+y Σ pa (W441)n4) V
DRTEE
Lemma 2
P (v)
27 7 TV(#(v)) m (va))
(V(((V)), T (v()) +y
Σ Uk JER
p= (TV- (Mom)))
JLEY
Lamma 3
P (wn (w) [T((((()
2
( (gx)) 1+ 7 (d (d') (p' (-1ame), p= P= ((1ma))
+y 1-
1-
PK (2ulv)
11
(TV(T (5 n T (yw)) + Wd' (-1ay,), pn (+184))
1-
Σ BEY
P (wnlv)
1+7 d' di (5 vk3
1-7
16

---

Published as a conference paper at ICLR 2021
Thus, by induction, it follows that for all n:
1 1+ TV - -
d" (Jy
which completes the proof.
C BISIMULATION METRICS
Notation. We the notation defined in Section 2.
Bisimulation metrics (Givan et al., 2003; Ferns et al., 2011) define a pseudometric d. S x S where d.d (xr is defined in terms of distance between immediate rewards and next state distributions. Define FC M > M by
FS)d( max R(x, a) R(v. a)l + (W( (P*()) P-(lv) aεA
(C.1)
then, FL has unique fixed point d. which is bisimulation metric. Fi the 1-Wasserstein metric W1 : MI + Mp. The J-Wasserstein distance (d) under the pseudometric d can be computed using the dual linear program:
Σ max u.V
P(z)u=
Σ yEy
Ply)u, subject to va N,yEy uz - d(z,v) v)
ZEX
Since we are only interested in computing the coupling between the states in X and V, the above formulation assumes that Px(v) for all y and Py(x)) 0 for all x X'. The computation of d-, is expensive and requires a tabular representation of the states, rendering it impractical for large state spaces. On-policy bisimulation (Castro, 2020) (e.g., x --bisimulltiion tied to specific behavior policies and is much easier to approximate than bisimulation.
D POLICY SIMILARITY METRIC
D.1 COMPUTING PSM
In general, PSM for given DIST across MDPs Ma and My is given by d' (x. y) DIST ( + 7W (d' ( (31x),
(D.1)
Since our main focus showing the utility of PSM for generalization, we simply use environments where PSM can be computed using dynamic programming. Using similar observation to Castro (2020), we assert that the recursion for d' takes the following form in deterministic environments:
d" [x, 3) = DIST (Tx (z), =5())) + 2dd"(a',3) (a',3/).
(D.2)
where x' x' (x) (y) are the next states from taking actions x (x), क. (t) from states x, respectively. Furthermore, we assume that DIST between terminal states from Mxx and My is zero. Note that the form of Equation D.2 closely resembles the update rule in Q-learning and as such, can be efficiently computed with samples using approximate dynamic programming. Given access to optimal trajectories x = { and T5 {ve J where 21+1 Px (Ix) and yr+1 = PT (y²). Equation D.2 can solved using exact dynamic programming: we provide pseudocode Section J.1.
There are other ways to approximate the Wasserstein distance in bisimulation metrics (e.g.. Ferns al., 2006; 2011; Castro, 2020; Zhang et al., 2021). That said, approximating bisimulation (or PSM) for stochastic environments remains an exciting research direction (Castro, 2020). Investigating other distance metrics for long-term behavior difference in PSM also interesting for future work.
D.2 PSM CONNECTIONS TO DATA AUGMENTATION AND BISIMULATION
Connection to bisimulation. Although bisimulation metrics have appealing properties such as bounding value function differences (e.g., (Ferns & Precup, 2014)), they rely on reward information and may provide meaningful notion of behavioral similarity in certain environments. Propo- sition D.1 implies that states similar under PSM would have similar optimal policies yet can have arbitrarily large bisimulation distance between them.
17

---

Published as a conference paper at ICLR 2021
20
yo
do, Tx "Cake"
C1 "No Cake"
dos Ty Caaee"
al "No Cake"
0g: "Cake"
ag- "Cake"
d0
00s
122
1/1
1/2
a1
aD
d1 "No Cake"
01 Ty "No Cake"
Figure D.I: Cyan edges represent actions with a positive reward, which also the optimal actions. Zero rewards everywhere else. z0. Yo are start states while #2. M2 are the terminal states.
Proposition D.I. There exists environments Mx and My such that V(r,y) L where L ((r,y) x E X, y E y, d' (x, = 0}. d=, Rmaax Rminn
- for any given e > 0.
For example, consider the two semantically equivalent environments Figure D.1 with π (mo) TH (yo) dy and x (x1) x
but different rewards rz: respectively. Whenever ry > + 1/4) Tz. bisimulation metrics incorrectly imply that 2o is more behaviorally similar to yi
than y/-
For the MDPs shown in Figure D. to determine which y state behaviorally equivalent to To. we look at the distances computed by bisimulation metric d. and x -i-bimiuulation metric d=
d- yo) = d' (ro. yo) = (1 7) ryl (20, y1) max ((1 + Fz: rg), d= (=o, v1) Ir, -rzl Yzz Thus, r > + 1/7) rx implies that ( 11) d. (1o: yo) as well (=.. < d=, (30: Bo)-
Connection to data augmentation. Data augmentation often assumes access optimality invariant transformations, e.g., random crops or flips in image-based benchmarks (Laskin et al., 2020a; Kostrikov al., 2020). However, for certain RL tasks, such augmentations can erroneously alias states with different optimal behavior and hurt generalization. For example, the image observation is flipped in goal reaching task with left and right actions, the optimal actions would also be flipped to take left actions instead of right and vice versa. Proposition D.2 states that PSMs can precisely quantify the invariance of such augmentations.
Proposition D.2. For an MDP Mx and its transformed version M the data augmentation db. d' (x, (z)) indicates the optimality invariance of for any x E X.
D.3 PSM WITH APPROXIMATELY- POLICIES
Generalized Policy Similarity Metric for arbitrary policies. For a given DIST, we define generalized PSM d (S II) x (S II) - R where is the set of all policies over S. satisfies the recursive equation:
(D.3) Since DIST is assumed to be a pseudometric and Wi is probability metric, it implies that d pseudometric as (1) d is non-negative, that is, d((x, 11), T2)) 2 0, (2) d is symmetric, that is, d((x, m1). (v. R2)) = d((x, m1), (y. T2)), and d satisfies the triangle inequality, that is, d((2, m1), (y, T2)) < d((z, mI), Ts)) + d((2, nT), T2)).
d((x, æ1), (y, m2)) = DIST (x), =e(w)) + 7WV (d) (P p*>(,))))
Using this notion of generalized PSM, we show that the approximation eтror in PSM from using a suboptimal policy bounded by the policy's suboptimality. Thus, for policies with decreasing suboptimality, the PSM approximation becomes more accurate, resulting in improved PSEs.
Proposition D.3. [Approximation error in PSM] Let d S x S + R be the approximate PSM computed using suboptimal policy Ω defined over S, that is, d(x,y) = (IsT( - Willd) (-1xx, p* (-lv)) We have:
ld" (z. v) - d(.,, v) <
(1,,,)),(
d(ly, (U, (1,
Long-term suboptimality difference from x
Long-term suboptimality difference from y
18

---

Published as a conference paper at ICLR 2021
Proof. The PSM d' and approximate PSM d are instantiations of the generalized PSM (Equation D.3) with both input policies as n² and respectively.
(xxx)) = ((0, x (v, )) < d((,, (d),
+ d((2, (v. a) :: d(a,y) v) = d((e, R), (v.
d'(),)) - d(z, y) ((,,))); (d (v. Similarly, d(z,y) d' (z, v) d(<z, d (v. 8), (U,
E L2 METRIC EMBEDDINGS
Another common choice (Zhang et al., 2021) for learning metric embeddings is to use the squared loss (i.e., 12-loss) for matching the euclidean distance between the representations of pair of states to the metric distance between those states. Concretely, for given d' and representation fe, the loss L()) E...,, (s.)
fo(ss)ll2 d' (st, sy)) is minimized. However, might be too restrictive to match the exact metric distances, which we demonstrate empirically by comparing (2
metric embeddings with CMEs (Section 5.2).
F EXTENDED RELATED WORK
Generalization across different tasks used to be described as transfer learning. In the past, most transfer learning approaches relied on fixed representations and tackled different problem formulations (e.g., assuming shared state space). Taylor & Stone (2009) present comprehensive survey of the techniques at the time, before representation learning became so prevalent in RL. Recently, the problem of performing well in different, but related task, started be seen problem of generalization; with the community highlighting that deep RL agents tend to overfit to the environments they are trained on (Cobbe al., 2019; Witty al., 2018; Farebrother et al., 2018; Juliani et al., 2019; Kostrikov et al., 2020; Song et al., 2019; Justesen et al., 2018; Packer et al., 2018).
Prior generalization approaches are typically adapted from supervised learning, including data aug- mentation, regularization (Cobbe et al., 2019; Farebrother et al., 2018) 8), stochasticity (Zhang et al., 2018c), noise injection (Igl al., 2019; Zhang et al., 2018a), more diverse training conditions (Ra- jeswaran et al., 2017; Witty et al., 2018) and self-attention architectures (Tang et al., 2020). In contrast, PSEs exploits behavior similarity (Section 3), a property related to the sequential aspect of RL. Fur- thermore, for certain RL tasks, it can be unclear what an optimality invariant data augmentation would look like (Section 5.3). PSM can quantify the invariance of such augmentations (Proposition D.2).
Meta-leamnig is also related to generalization. Meta-learning methods try to find parametrization that requires a small number gradient steps to achieve good performance on a new task (Finn et al., 2017). In this context, various meta-learning approaches capable of zero-shot generalization have been proposed (Li et al., 2018; Agarwal et al., 2019; Balaji et al., 2018). These approaches typically consist in minimizing the loss in the environments the agent is while adding auxiliary loss for ensuring improvement in the other (validation) environments available to the agent. Nevertheless, Tachet des Combes et al. (2018) has shown that meta-learning approaches fail the jumping task which we also observed empirically. Others have also reported similar findings (e.g., Farebrother et al., 2018).
There are several other approaches for tackling zero-shot generalization in RL, but they often rely on domain-spenn speciffc information. Some examples include knowledge about equivalences between entities in the environment (Oh et al., 2017) and about what is under the agent's 's control (Ye et al., 2020). Causality-basse methods are a different way of tackling generalization, but current solutions do not scale to nigh-dimensiooal observation spaces (e.g., Killian et al., 2017; Perez et al., 2020; Zhang et al., 2020).
19

---

Published as a conference paper at ICLR 2021
G JUMPING TASK WITH PIXELS
OBSTACLE POSITION- 20, FLOOR IEEGGH
OBSTACLE POSITION 30, FLOOR HerGHt
10
Figure G.1: Optimal trajectories on the jumping tasks for different environments. Note that the optimal trajectory sequence of right actions, followed hy single jump at a certain distance from obstacle, followed by right actions.
Detailed Task Description. The jumping task consists of agent trying to jump over an obstacle on floor. The environment is deterministic with the agent observing reward of +1 at cach time step. the agent successfully reaches the rightmost side of the screen, receives a bonus reward of + +100; if the agent touches the obstacle, the episode terminates. The observation space is the pixel representation the environment, as depicted in Figure 1. The agent has access to two actions: right and jump. The jump action moves the agent vertically and horizontally the right.
Architecture. The neural network used for Jumping Task experiment is adapted from the Nature DQN architecture. Specifically, the network consists of 3 convolutional layers of sizes 32, 64, 64 with filter sizes 8 x 8, 4 4 and 3 x and strides 4, 2, and 1, respectively. The output of the convnet fed into single fully connected layer of size 256 followed by 'ReLU' non-linearity. Finally, this FC layer output is fed into linear layer which computes the policy which outputs the probability of the jump and right actions.
Contrastive Embedding. For all our experiments, we use a single ReLU layer with À = 64 units for the non-linear projection obtain the embedding 29 (Figure A.I). We compute the embedding using the penultimate layer in the jumping task network. Hyperparameters are reported in Table G.2.
Total Loss. For jumpy world, the total loss is given by Si. + a.MEM where Ln, is the imitation learning loss, .CME is the auxiliary loss for learning PSEs with the coefficient a.
Test tasks solved
PSEs
tasks solved
PSEs
% Test tasks solved
PSEs
50%
Dropout + (2 reg.
50%
Dropout + reg.
50%
Dropout 12 reg.
30%
20%
20%
20%
10%
10%
10%
0%
0%
500
1000
1500
2000
500
1000
1500
2000
500
1000
1500
2000
Training Iterations
Training Iterations
Training Iterations
Figure G.3: Test performance curves in the setting without data augmentation on the "wide", "narrow", and random grids described in Figure We plot the median performance across 100 runs. Shaded regions show 25 75 percentiles.
20

---

Published as a conference paper at ICLR 2021
BISIMULATION
PSM
175
1.75
950
350
150
1.50
440
340
125
1.25
1/0
3/30
100
1.00
1/20
75
3/20
0.75
1/10
50
3/10
0.50
yo
25
3/0
0.25
20
3/0
20
30
3550
270
10
230
E40
350
0.00
(a) Bisimulation metric
(b) Policy Similarity metric
1.0
BISIMULATION
TRUE SIMILARITY
4000
3/0
3/50
0.8
940
3000
740
930
3/30
2000
Y20
3/20
0.4
3/00
1000
Y10
0.2
Yo
yo
10
T20
1:30
T40
2550
3'10
3220
3:00
344
350
0.0
(c) n" -Bisimulation metric
(d) True State Similarity
Figure G.2: PSM vs. Bisimulation. Visualizing PSM and bisimulation metrics with discount factor 7 = 0.99. and y correspond to the states visited by optimal policies in the two environments with obstacle at positions 25 and 45, respectively. For each grid, the (i,j)th location shows the distance assigned the metric to the states 2; and リ,- Lower distances (lighter shades) imply higher similarity between two states. The bottom right figure shows which states are equivalent to each other. The ranges are different for each metric as bisimulation metrics utilize reward differences while PSM uses the total variation (TV) distance between policies. Note that the large bisimulation distances due to the fact that the reward at the terminal state is set 100.
Test tasks solved
100%
% Test Test tasks solved
100%
RandConv PSEs
% Test tasks solved
100%
RandConv
80%
80%
80%
60%
60%
60%
40%
RandConv PSEs
RandConv PSEs RandConv
RandConv
500
1000
1500
2000
1000
1500
2000
1000 2000
Training Iterations
Training Iterations
Training Iterations
Figure G.4: Test performance curves in the setting with data augmentation on the "wide", "narrow", and random grids described in Figure We plot the median performance across 100 runs, Shaded regions show 25 and 75 percentiles.
21

---

Published as a conference paper at ICLR 2021
G.1 JUMPING TASK FROM COLORS
Table G.I: Percentage (%) of test tasks solved when trained the "wide" grid with both red and green obstacles. The numbers we report are averaged across 100 runs. Standard error reported between parentheses.
Method
Red (%)
Green (%)
RandConv
6.2 (0.4)
99.6 (0.2)
Figure G:5: Jumping task with colored ob- stacles. The agent needs to jump over the red obstacle strike the green obstacle.
Dropout and le reg.
19.5 (0.2)
100.0 (0.0)
RandConv + PSEs
29.8 (1.3)
99.6 (0.2)
PSEs
37.9 (1.9)
100.0 (0.0)
INPUT
AUCMENTATION
AUGMENTATION
AUGMENTATION
AUGMENTATION
AUGMENTATION
INPUT
AUGMENTATION
AUGMENTATION
AUGMENTATION
AUGMENTATION
AUGMENTATION
Figure G.6: Randcony enforces color invariance. The first column shows the original observation where the top corresponds to the observation from task with green obstacle jumping task while the bottom row corresponds the red obstacle jumping task. Columns 2-6 in each show augmentations by applying RandConv. The augmentations show that RandConv tries encode invariance with respect to obstacle color.
G.2 HYPERPARAMETERS
For hyperparameter selection, we evaluate agents on a validation set containing 54 unseen tasks in the "wide" grid and pick the parameters with the best validation performance. The validation set (Figure G.7) was selected by using environments nearby to the training environments whose floor height differ by 1 or whose obstacle position differ by 1.
Table G.2: Common hyperparameters across all methods for all jumping task experiments.
v
V
Hyperparameter
Value
Learning rate decay
0.999
Training epochs
2000
Optimizer
Adam
Batch size (Imitation)
256
Num training tasks
18
P-scale Parameter (B)
0.01
Embedding size (k)
64
Figure G.7: Unseen 56 validation tasks, labeled with V, used for hyperparameter selection.
Batch Size (LCMF)
57
57
22

---

Published as a conference paper at ICLR 2021
Table G.3: Optimal hyperparameters for reporting results in Table These hyperparameters selected using the "wide" grid by maximizing final performance on validation set containing 56 unseen tasks. All grid configurations in Table use these hyperparameters.
Hyperparameter
Dropout and lg-reg.
PSEs
RandConv RandConv + PSEs
Learning Rate
4x 4x10-3
3.2x 10-3
7x10-3
2.6x10-3
f2-reg. coefficient
10-4
1x10-5
Dropout coefficient
3x10-1 Auxiliary loss coefficient (ax) _
Contrastive Temperature (1/X)
1.0
5x10-1
1x10
5.0
Table G.4; Optimal hyperparameters for reporting results in Figure 5.3. These hyperparameters are selected using the "wide" grid by maximizing final performance on a validation set containing 56 unseen tasks.
Hyperparameter
Dropout and fg-reg.
PSEs
RandConv RandConv PSEs
Learning Rate
4x 4x10-3
6x10-
5x10-3
2.6x
fz-reg. coefficient
4.3x10-*
7x10-5
Dropout coefficient
3x 3x10-1
Contrastive Temperature (1/A)
5x10-1
5x 6x10-1
Auxiliary loss coefficient (a)
5.0
5.0
Table G.S: Optimal hyperparameters for reporting ablation results Table 2. These hyperparameters are selected using the "wide" grid by maximizing final performance on validation set containing 56 unseen tasks.
PSM
π bisimulatiin
Hyperparameter
CMEs
l ee-meeddinngs
CMEs
fg-embeddings
Learning Rate
4x 1x10-3
5x10-1
4.7x10-4
1x10-1
Contrastive Temperature (L/A)
1.0
5x10-1
Auxiliary loss coefficient (a:)
5.0
1x10-1
1x10-1
1x10-6 6
Please note that Table G.3 and Table G.4 correspond to two different tasks: one uses the standard jumping task with white obstacles, while other uses colored obstacles where the optimal policies depend on color. For fair comparison, we tune hyperparameters for all the methods using Bayesian optimization (Golovin et al., 2017). We use the best parameters among these tuned hyperparameters and the ones found in Table G.3, leading different parameters for both PSEs as well as RandConv. Evaluating PSEs with the jumping task hyperparameters from Table G.3 instead of the ones in Table G.4 leads to small drop (-4%) on the jumping task with colors (Section 5.3). Nevertheless, PSEs still outperform other methods in Section 5.3.
H LQR: ADDITIONAL DETAILS
Optimal control with linear dynamics and quadratic cost, commonly known as LQR, has been increasingly used simplified surrogate for deep RL problems (Recht, 2019). Following Song et al. (2019); Sonar et al. (2020), we analyze the following LQR problem for assessing generalization:
minimize
EngND Eee T J.Qs: + a. Raj
[0.1 W. Wa
(H.1)
subject to 8111 = Ast + Bat, Of =
8t. Gt = Kor,
where is the initial state distribution, 8t E R is the (hidden) true state at time t. a, E Re is the control action, and K is the linear policy matrix. The agent receives the input observation ot. which is a linear transformation the state s,, W, and Wa are semi-orthogonal matrices to prevent information loss for predicting optimal actions. An environment corresponds particular choice
23

---

Published as a conference paper at ICLR 2021
Table H.l: LQR generalization performance: Absolute emror in LQR cost, w.L.t the oracle solver (which has access true state), various methods trained with na distractors on environments. The reported mean and standard deviations are across 100 different seeds. Lower error is better.
Number of Distractors (nz)
Method
500
1000
10000
Overparametrizaaioon (Song et al., 2019)
25.8 (1.5)
24.9 (1.1)
24.9 (0.4)
IPO (Sonar al., 2020) (IRM + Policy opt.)
32.6 (5.0)
(2.8
24.8 (0.4)
Weight Sparsity (f1-reg.) PSM (State aggregation)
28.2 (0.0)
(0.0)
28.2 (0.0)
0.03 (0.0)
0.03 (0.0) 0.02 (0.0)
of Wa; all other system parameters (A, B, Q, R, Wc) are fixed matrices which are shared across environments and unknown to the agent. The agent learns the policy matrix K using N training environments based on Equation H.1. test time, the learned policy is evaluated on environments with unseen Wa.
The generalization challenge in this setup is ignore the distractors: W.s, E R represents the state features invariant across environments while Was; RR is high-dimensional distractor of size ns Odd respectively, such that n, << Ad. Furthermore, the policy matrix which generalizes
10 W.P." T
across all environments is K2
where P. corresponds to the optimal LQR solution
0
with access to state st. However, for single environment with distractor Wa. multiple solutions
10a P.P
Va [0, 1]. Note that the distractors are order of magnitude larger than the invariant features in or and dependence on them is likely to cause the agent act erratically on inputs with unseen distractors, resulting in poor generalization.
exist, for instance, K. (1 - a:) W.P.T
We use overparametrized policies with two linear layers, i.e., K K2, where K1(0) is the learned representation for observation o. We learn K using gradient descent using the combined cost on 2 training environments with varying number of distractors. We aggregate observation pairs with near-zero PSM by matching their representations using squared loss. We use the open-source code released by Sonar et al. (2020) for our experiments.
Table H.2: An overview of hyper-parameters for LQR.
Parameter
Setting
A
Orthogonal matrix, scaled 0.8
B
T20x20
20
na
20
Q
720 20
R
T20
K₁ Vi
Orthogonal Initialization, scaled 0.001
Wa
Random semi cemi-orthogonal matrix
The reliance on distractors for IPO also highlights a limitation of IRM: if a model can achieve solution with zero training error, then any such solution is acceptable by IRM regardless of its generalization ability - common scenario with overparametrized deep neural nets (Jin et al., 2020).
H.1 NEAR-OPTIMALITY OF PSM AGGREGATION
Conjecture 1. Assuming zero state aggregation error with policy similarity metric (PSM), the policy matrix K learned using gradient descent is independent of the distractors.
Proof. For LQR domains x, y, an observation pair (o, ", or") has zero PSM iff the underlying state s; same for both the observations in the pair. This is true, (a) both domains has the same transition
24

---

Published as a conference paper at ICLR 2021
dynamics, specified by Equation H.1, and (b) the optimal policy is deterministic and is completely determined by the current state 8 at any time
0.1 We
Wa= 5, and o," 0..1 We
Assume of
St for distractor semi-orthogonal matrices Wax and
de
Way, respectively. Furthermore, the representation is given by Ki (o, (o,*) and Ki (o,") respectively. Assume that K1 [Ks Kd] where K, /2/xa, and Ka E Rhxad and K₁ E (gxx(n,tna)
Zero stale-aggregation error with squared loss implies that for pair (o,", or") corresponding to 8t.
0
K1(0, - - - = K1 Wa Way = Ka(Wa= W )s; =0
(H.2)
As Equation H.2 holds for all states visited by the optimal policy in an infinite horizon LQR, follows that Ka(Wa a- W 0.
Furthermore, it is well-known that gradient descent tends to find low-rank solutions due to implicit regularization (Arora et al., 2019; Gunasekar et al., 2017), e.g., with small enough step sizes and initialization close enough to the origin, gradient descent on matrix factorization converges to the minimum nuclear norm solution for 2 layer linear networks (Gunasekar et al., 2017). Based on this, we conjecture that K₃ = 0 which we found to be true practice.
I DISTRACTING CONTROL SUITE
We use the same setup as Kostrikov al. (2020); Stone et al. (2021) for implementation details and training protocol. For completeness, we describe the details below.
Dynamic Background Distractions. In Distracting Control Suite (Stone et al., 2021), random backgrounds are projected from scenes of the DAVIS 2017 dataset (Pont-Tuueet et al., 2017) onto the skybox of the scene. To make these backgrounds visible for all tasks and views, the floor grid is cemi-transparent with transparency coefficient 0.3. We take the first 2 videos in the DAVIS 2017 training set and randomly sample a scene and a frame from those at the start of every episode. In the dynamic setting, the video plays forwards or backwards until the last or first frame is reached which point the video is played backwards. This way, the background motion is always smooth and without "cuts".
Soft Actor-Critic Soft Actor-Critic (SAC) (Haarnoja et al., 2018) learns a state-action value function Qo, stochastic policy my and temperature to find an optimal policy for an MDP (S, p,,, 7) by optimizing ryddiscountd maximum-entrow entropy objective. 0 is used generically to denote the parameters updated through training in each part of the model. The actor policy (e(a (s)) is parametric tanh- Gaussian that given samples parametric mean and standard deviation.
tanh(Ho(s;) + (o(ss)c). where ~ NO, 1) and pi and dg are
The policy evaluation step learns the critic Qo(si, a.) network by optimizing single-step of the soft Bellman residual
(D) (D) = E =l a:) ) yt)
d(~x(-1=))
Y: r(st, 0z) 7Q (8), o4) a log To(a)(s())
where D replay buffer of transitions. 0' an exponential moving average of the weights. SAC uses clipped double-Q learning, which we omit for simplicity but employ in practice.
The policy improvement step then fits the actor policy ne (a.ls:) network by optimizing the objective
(D) = Es,nD! DKL (TO (-1s:) exp Qo(ss. ))]]
a
Finally, the temperature is learned with the loss
(D) F,
(-a log To (a - aFH),
S a a.
where H. IR is the target entropy hyper- parameee that the policy tries to match, which in practice is usually set to H -l
25

---

Published as a conference paper at ICLR 2021
Figure I.1: DCS Test Environments: Snapshots of environments used for evaluating generalization on
Distracting Control Suite. Random backgrounds are projected from scenes the first videos of DAVIS 2017 validation dataset.
26

---

Published as a conference paper at ICLR 2021
1.1 ACTOR AND CRITIC NETWORKS
clipped double Q-learning for the critic, where each Q-function is parametrized as a 3-layer MLP with ReLU activations after each layer except of the last. The actor is also 3-layer MLP with ReLUs that outputs mean and covariance for the diagonal Gaussian that represents the policy. The hidden dimension is set to 1024 for both the critic and actor.
Following Kostrikov et al. (2020),
1.2 ENCODER NETWORK
We employ the encoder architecture from Kostrikov et al. (2020). This encoder consists of four convolution layers with 3 x 3 kernels and 32 channels. The ReLU activation is applied after each convolutional layer. We use stride to everywhere, except of the first convolutional layer, which has stride 2. The output of the convnet is feed into single fully-connected layer normalized by LayerNorm. Finally, we apply tanh nonlinearity to the 50 dimensional output of the fully- connected layer. We initialize the weight matrix of fully-connected and convolutional layers with the orthogonal initialization and set the bias to be zero. The actor and critic networks both have separate encoders, although we share the weights of the conv layers between them. Furthermore, only the critic optimizer is allowed to update these weights (i.e., we stop the gradients from the actor before they propagate to the shared convolutional layers).
1.3 CONTRASTIVE METRIC EMBEDDING LOSS
For all our experiments, we use a single ReLU layer with A = 256 units for the non-linear projection to obtain the embedding 29 (Figure A. We compute the embedding using the penultimate layer in the actor network. For picking the hyperparameters, we used 3 temperatures [0.1, 0.01, 1.0] and auxiliary LCME loss coefficients 31, 10] using "Ball In Cup Catch" the validation environment. All other hyperparameters are the same as prior work (see Table I.2).
We approximate optimal policies with the policies obtained after training DrQ agent for 500K environment steps. Since a given action sequence from this approximate policy has the same perfor- mance across different training environments, we compute the PSM across training environments, via dynamic programming (see Section J.1 for pscudo-code), using such action sequences.
Total Loss. The total loss is given by LRI. + aLME where LR. is the reinforcement learning loss which combines Jx (D), Jn (D), and Jn (D)), while LEME is the auxiliary loss for learning PSEs with the coefficient a.
1.4 TRAINING AND EVALUATION SETUP
For evaluation, we the first 30 videos from the DAVIS 2017 validation dataset (see Figure I.1). Each checkpoint is evaluated by computing the average episode return over 100 episodes from the unseen environments. All experiments are performed with five random seeds per task used compute means and standard deviations/errors of their evaluations. We use K = 2, MM = 2 as prescribed by Kostrikov et al. (2020) for DrQ. Following Kostrikov et (2020) and Stone et al. (2021), we use different action repeat hyper-parameter for each task, which we summarize in Table I.3. We construct an observational input as a 3-stack of consecutive frames (Kostrikov et al., 2020), where each frame is an RGB rendering of size 84 x 84 from the Oth camera. We then divide each pixel by 255 to scale it down to [0, 1] range. For data augmentation, we maintain temporal consistency by using the same crop augmentation across consecutive frames.
Table I.1: Optimal hyperparameters for PSE auxiliary loss for reporting results in Table 3.
Hyperparameter
Setting
Contrastive temperature (1/
0.1
Auxiliary loss coefficient (a)
1.0
T-scale parameter (3)
0.1
Batch Size (LCME)
128
1000 // Action Repeat
27

---

Published as a conference paper at ICLR 2021
Table 1.2: syyper-aaammters taken from Kostrikov et al. (2020) in the Distracting Control Suite experiments.
Parameter
Setting
Replay buffer capacity
100, 000
Seed steps
1, 1,000
Batch size (DrQ)
512
Discount y
0.99
Optimizer
Adam
Learning rate
10
Critic target update frequency
2
Critic Q-function soft-update rate
0.01
Actor update frequency
2
Actor log stddev bounds
[- 10,
Init temperature
0.1
Table L3:
The action repeat hyper-parameter used for each task in the Distracting Control Suite benchmark.
Task name
Action repeat
Cartpole Swingup
8
Reacher Easy
4
Cheetah Run
4
Finger Spin
2
Ball In Cup Catch
4
Walker Walk
2
1.5 GENERALIZATION CURVES
DrQ PSEs
DrQ
ball in catch
cartpole-swingup
cheetah-run
800
Episode Return
800
300
600
600
200
400
400
200
100
200
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.1
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Environment Steps (x10*)
Environment Steps (x10*)
Environment Steps (x10')
finger-spin
reacher-easy
walker-walk
1000
Episode Return
750
500
600
500
250
400
250
200
0
0.0
0.1
0.2
0.3
0.4
0.0
0.1
0.2
0.3
0.4
0.0
0.2
0.3
0.5
Environment Steps (x10*)
Environment Steps (x10*)
Environment Steps (x 1x'10))
Figure 1.2: Random Initialization. Generalization performance on unseen environments over the course of training. The agent is initialized randomly. DrQ augmented with PSEs outperforms DrQ. We plot the average episode return across 5 seeds and the shaded region shows the standard deviation. Each checkpoint is evaluated using 100 episodes with unseen backgrounds
28

---

Published as a conference paper at ICLR 2021
DrQ + PSEs (Pretrained)
DrQ (Pretrained)
ball in cup-catch
cartpole-swingus
cheetah-run
Episode Return
800
800
300
600
200
400
100
0.0
0.1
0.2
0.3
0.4
0.0
0.1
0.2
0.3
0.1
0.0
0.1
0.2
0.3
0.5
Environment Steps (x10*)
Environment Steps (x10*)
Environment Steps (x100)
finger-spin
reacher-easy
walker-walk
Episode Return
800
800
600
600
400
0.0
0.1
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.0
0.2
0.3
0.5
Environment Steps (x10%)
Environment Steps (x10")
Environment Steps (x 10*)
Figure 1.3: Pretrained Initialization. Generalization performance on unseen environments over the course of training. The agent is initialized using pretrained DrQ agent. DrQ augmented with PSEs outperforms DrQ on most of the environments. We plot the average return across 5 seeds and the shaded region shows the standard deviation. Each checkpoint is evaluated using 100 episodes with unseen backgrounds.
J PSEUDO CODE
J.1 DYNAMIC PROGRAMMING FOR COMPUTING PSM
def cetrri_ffxed_pee
matriai gamma-0. 99, eps-le-7):
"""DP for calculating PSM in environments with deterministic dynamics.
Args:
cost_matrix: matrix where entries at index (i, 3) is (IST y-j)
gamma: Metric discount facter.
eps: Threshold for stopping the fixed point iteration.
mp. zeoos_ssslike like (cost_matrix) def (d
10
11
dneew = 1 cost_matrix
discounted_d d cur gamma d_cur
13
d_new[:-l, :-1]
di scounted_d_ccur [1:, 1:]
14
d_new[:-] -1] +- di scounted_d_cou [1;, -1]
15
d_new[-1, :-1] +- di poounted_d_curr [-1, 1:]
return new
18
while True:
19
d_new operator (d)
20
if up.sum(np.abl (d d_new) < eps:
break
22
else:
23
d__ew[[]]
24
return d
J.2 CONTRASTIVE LOss
def contrastive_les cimilarity... matriix,
metric carri__alleee
temperature,
29

---

Published as a conference paper at ICLR 2021
beta=1.0)
**nnttraiiv Loss with embedding similarity "**
metric_shape shape (metric_values)
## 2_Ntheta (X) embedding_1
mwwmodde.reprcceetttateeeee (X)
# z_\theta (Y) embedding_2 nn_model. represtatino (Y)
similarity_matrix cosine_similarity (embeddigg,, erbedding_2
10
# metric_values (X, Y)
11
similarity.matrik /= temperature
12
neg_iogitsl similarity_matris
0ol col_indices tf. (tf..argmi argmin (metric_values, axis-l), dtype-tf.int32) pos_ _es_indieesl tf.stack(
15
16
range (metric_shape [0], [0], (type=tf.int322) col_indices axis-1)
pos_logitsl = tf.gather nd (similarity.
pos_ indicesl)
19
netric_values /= beta
20
similarity_measure = tf. oxp (xttriitiicccaalleee
pos_wcights1 = gtf.aathr__nd (metriccaaet pos_indicosi)
pos_logitsl pos_ pos_weightss
begative_weights.
(11. similariity_measure) 1e-8)
24
neg_logital +- tf. tf_ttensor_ccaatee_____pe.ee
25
negative_weights, pos_indicesl, pos_weights!)
neg neg_Login Logitsl w.w.tth.....m reduce Logsumexp (neg logitsl, (neg_Logitsl,, axis=1) return tf.reduce mean
28
pos los_ttg # Equation 4
a
30


Reviews:
[{"review_id": "igTS3zvKhb5", "reviewer": "AnonReviewer2", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper defines a new metric for calculating similarities between states. The policy similarity metric is basically based on the $\\pi$-bisimulation metrics, but the l1 distance term of rewards is modified to the distance between states transitioned by the grounded policy. The experiments performed with several benchmarks, and the gain seems to be consistent.\n\n[Quality]\n\nThe paper is clearly written overall and well arranged.\n\n[Originality & Significance]\n\nThe originality and significance are sufficient for acceptance, as the authors devised a novel metric.\n\n[Strengths]\n+ Even though I didn’t thoroughly check all of the derivations, the metric seems to be more precise than $\\pi$-bisimulation, similarity between states originated from different environments.\n+ This approach can be easily applied to various RL models.\n\n[Weaknesses]\n- In Table 1, experiments with RandConv only contain RandConv and RandConv+PSEs. It should be better to compare to other methods, at least Bisimulation, to show how much the PSE model performs well. Not only Table 1, I feel overall comparison results are not sufficient.\n- The Distracting Control Suite dataset, which is the core dataset in the experiments, is not currently available, so verifying the results with this dataset is difficult.\n\n[Questions or Additional Comments]\n1. Figure 2 is not self-explanatory and hard to understand. I think adding axes or an additional explanation in the caption is required.\n2. Please elaborate on the reason for using the Distracting DM Control Suite dataset rather than other datasets.\n3. According to Tables G.3 and G.4, hyperparameters are quite different in each experiment setting. Does the proposed method’s performance highly depend on hyperparameter settings?\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "SotEr9ZSwq", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose an interesting direction toward generalization in reinforcement learning. The authors propose a policy similarity metric (PSM) for measuring behavioral similarity between states. States are considered similar if the optimal policies in those states (as well as future states) are similar. The authors present a contrastive representation procedure to embed state similarity metrics, which are instantiated with PSM to yield policy similarity embeddings (PSEs). \n\nI like the idea of looking at the similarity between policies by examining the difference in local and long-term behavior. Furthermore, I like that the contrastive metric embedding approach is relatively simple and that PSEs can be learned alongside RL using an auxiliary objective.\n\nQuestions and comments:\n1. What informed the train test split in the problem setup on page 4?\n2. I would like to see how this approach performs on environments specifically designed to test generalization, such as the ProcGen [1] and MineRL [2] environments.\n\n[1] Cobbe, et al. Leveraging procedural generation to benchmark reinforcement learning. arXiv:1912.01588, 2019.\n[2] Guss, et al. The MineRL competition on sample efficient reinforcement learning using human priors. arXiv:1904.10079, 2019.", "extra_scores": {}, "extra_text": {}}, {"review_id": "LdMS5eyWK-g", "reviewer": "AnonReviewer1", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Motivated by the issue of RL policy generalization, this paper explores improving generalization via a learned contrastive representation to embed states in, rather than through data augmentation or regularization alone. They demonstrate that their approach of policy similarity embeddings (PSEs) leads to improved generalization on several benchmarks, building on the method of policy bisimulation.\n\nStrengths:\n+ I think this approach is theoretically well-motivated, and the authors also give good intuition, especially in light of the latest work on self-supervised learning.\n+ Experimentally, I thought the controls were well-chosen and illustrated the utility of their method on each task they considered (Figures 4 and Tables 1 and 3).\n\nWeaknesses:\n-\tI would be curious to see how this approach works on more scaled up tasks with larger action spaces (Section 6.2 is in that direction, for instance), but the fact that it outperforms data augmentation techniques (including bisimulation transfer) in these simpler domains is promising.\n-\tIn the jumping task, no intuition is given as to why PSE underperforms in the “narrow grid” (Table 1, first row) compared to bisimulation transfer, and it would be helpful to explain the failure cases of their method in more detail generally.\n\nAs it stands, I think the ideas of this paper are interesting and novel, but I would like to see tasks with larger action spaces or more naturalistic inputs than the ones considered here, and in the cases where the method does not perform well, documentation and potential intuition for why that might be. Therefore, I recommend a weak accept.", "extra_scores": {}, "extra_text": {}}, {"review_id": "75NVC2Mljop", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \n\nThis work proposes a policy similarity metric (PSM) which differs slightly from the \\pi-bisimulation metric: while the \\pi-bisimulation metric defines the distance between two states in terms of the difference between the expected rewards obtained by following policy \\pi, the proposed policy similarity metric measures distance based on the similarity between the optimal policies at each state. The policy similarity metric can then be used to learn policy similarity embeddings (PSEs) via contrastive representation learning. These embeddings should in principle improve generalization in RL. \n\n\n\nStrengths: \n\n- The proposed policy similarity metric is well motivated for generalization in RL, which is an important problem. I also appreciate the theoretical upper bound on the suboptimality of transferring policies to a new environment.\n\n- The experiments are thorough and explore very diverse tasks. They highlight how PSEs can be used in imitation learning, LQR under spurious correlations, and RL settings; and can be combined with existing data augmentation techniques, including RandConv and DrQ. \n\n- The ablation studies and visualizations demonstrate, both quantitatively and qualitatively, the individual importance of the policy similarity metric and the contrastive metric embeddings.\n\n- The paper is very clearly written and well structured.\n\n\n\nWeaknesses:\n\n- The main weakness of the proposed policy similarity metric is that its computation requires optimal trajectories or approximately optimal trajectories (obtained, for example, through RL) from at least two training MDPs.\n\n- Building upon the previous point, I am curious about the utility of PSEs learned from increasingly suboptimal policies. I think including such a study would greatly help understand the sensitivity of learning PSEs to quality of the policies.\n\n- In settings like the jumping task from pixels, because of the reliance of the PSM on positive examples, data augmentation is required for good performance. \n\n\n\nRecommendation:\n\nI am recommending to accept this paper. I think the work is well motivated, approaches an important problem, and conducts very thorough experiments. I would still like to understand how useful PSEs are when learned from increasingly suboptimal policies.\n\n\n\nQuestions:\n\n- For the Distracting Control Suite tasks, are the “DrQ + PSEs” policies trained with both DrQ and PSEs or only PSEs? Specifically, I’m trying to understand if the gains are only from PSEs as opposed to the combination of PSEs and data augmentation.\n\n\n\nUpdates:\n\nThank you for the clarifications, new study, and theoretical justification. The new results demonstrate the utility of PSEs even when policies are not exactly optimal. My concerns were thoroughly addressed by the authors’ response, and I’ve updated my score accordingly.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper describes a new technique to generalize across different environments.  More precisely a new state similarity metric is defined with a contrastive learning embedding technique.  Unlike previous works that extend supervised learning techniques such as data augmentation and regularization to RL, the proposed approach takes into account the sequential nature of RL.    The reviewers unanimously praised the work in terms of theory, algorithm and empirical evaluation.  This is a novel and technically deep contribution that advances the state of the art for RL generalization.
Decision: Accept (Spotlight)