Cross-Trajectory Representation Learning for Zero-Shot Generalization in RL
Authors: Bogdan Mazoure
Ahmed M Ahmed
R Devon Hjelm
Andrey Kolobov
Patrick MacAlpine
Venue: ICLR 2022 Poster (2022)
Keywords: reinforcement learning
representation learning
self-supervised learning
procgen

Abstract: A highly desirable property of a reinforcement learning (RL) agent -- and a major difficulty for deep RL approaches -- is the ability to generalize policies learned on a few tasks over a high-dimensional observation space to similar tasks not seen during training. Many promising approaches to this challenge consider RL as a process of training two functions simultaneously: a complex nonlinear encoder that maps high-dimensional observations to a latent representation space, and a simple linear policy over this space. We posit that a superior encoder for zero-shot generalization in RL can be trained by using solely an auxiliary SSL objective if the training process encourages the encoder to map behaviorally similar observations to similar representations, as reward-based signal can cause overfitting in the encoder (Raileanu et al., 2021). We propose Cross-Trajectory Representation Learning (CTRL), a method that runs within an RL agent and conditions its encoder to recognize behavioral similarity in observations by applying a novel SSL objective to pairs of trajectories from the agent's policies. CTRL can be viewed as having the same effect as inducing a pseudo-bisimulation metric but, crucially, avoids the use of rewards and associated overfitting risks. Our experiments ablate various components of CTRL and demonstrate that in combination with PPO it achieves better generalization performance on the challenging Procgen benchmark suite (Cobbe et al., 2020).  

Published as a conference paper at ICLR 2022
CROSS- REPRESENTATION LEARNING FOR ZERO-SHOT GENERALIZATION IN RL
Bogdan Mazoure*1
Ahmed M. Ahmed*t attneddh@ttafoore edu Stanford University
bogdan bogdan..azourz mazoure aasourremmall. mogill. ca McGill University, Quebec AI Institute
Patrick MacAlpine
R Devon Hjelm
patrick. maca wwweesonnnnnnnm..o com Sony Al
devor hije mrjelmemiie crosoft com Université de Montréal, Quebec Al Institute, Microsoft Research
Andrey Kolobov
akol akolobov@microse com Microsoft Research
ABSTRACT
A highly desirable property of reinforcement learning (RL) agent - and a major difficulty for deep RL approaches - is the ability generalize policies Iearned on a few tasks over a high-dimensional observation space to similar tasks not seen during training. Many promising approaches to this challenge consider RL as a process of training two functions simultaneously: a complex nonlinear encoder that maps high-dimensional observations to latent representation space, and simple linear policy over this space. We posit that a superior encoder for zero-shot generalization in RL be trained by using solely an auxiliary SSL objective if the training process encourages the encoder to map behaviorally similar observations similar representations,
reward-based signal can cause overfitting in the encoder (Raileanu and Fergus, 2021). We propose Cross Trajectory Representation Learning (CTRL), method that runs within RL agent and conditions its encoder to recogmize behavioral similarity in observations by applying novel SSL objective pairs of trajectories from the agent's 's policies. CTRL can be viewed as having the same effect as inducing pseudo- bisimulation metric but, crucially, avoids the use of rewards and associated overfitting risks. Our experiments ablate various components of CTRL and demonstrate that in combination with PPO it achieves better generalization performance on the challenging Proegen benchmark suite (Cobbe et al., 2020).
INTRODUCTION
Deep reinforcement learning (RL) has emerged as a powerful tool for building decision-making agents for domains with high-dimensional observation spaces, such video games (Mnih et al., 2015), robotic manipulation (Levine et al., 2016), and autonomous driving (Kendall et al., 2019). However, while deep RL agents may excel the specific task variations they are trained on, learning behaviors that generalize across a large family of similar tasks, such as handling a variety of objects with robotic manipulator, driving under a variety conditions, or coping with different levels in a game, remains a challenge. This problem is especially acute in zero-shot generalization (ZSG) settings, where only a few sequential tasks are available to learn policies that are meant to perform well on different yet related tasks without further parameter adaptation. ZSG settings highlight the fact that generalization often cannot be solved by more training. as can be too expensive or impossible to instantiate all possible real-world deployment scenarios a-priori.
"Equal contribution. 'The author did part of the work for this paper while at Microsoft. Code link: https: / / thub. com hwmm/maa//eeeocoooocmmaa/amaaee/ www.w/ctrreeeree ic

---

Published as a conference paper at ICLR 2022
In this work, aim to improve ZSG in RL by proposing a new way of training the agent's representation, a now-dimensional summary of information relevant to decision makinng extracted from the agent's 's sighddmmessionll observations. Outside of RL, representation learning with ZSG, e.g. using unsupervised learning to obtain a representation that readily transfers to unseen classes in vision tasks (Bucher et al., 2017; Sylvain al., 2020; Wu et al., 2020). In RI., unsupervised representation learning in the form of auxiliary objectives can be used to provide a richer learning task-specific
help
signal over learning from reward alone, which helps the agent avoid overfitting information (Raileanu and Fergus, 2021). However, to our knowledge, no unsupervised learning method used in this way in RL has thus far been shown substantially improve performance in ZSG over end-to-end reward-bacc basd methods (e.g.. Cobbe et al., 2020).
We posit that using unsupervised (reward-free) learning find representations that capture behav- ioral similarity across different trajectories will improve ZSG in RL. We note that the bisimulation framework (Ferns et al., 2004) does this directly with rewards, optimizing an agent to treat states as behaviorally similar based on the expected reward, and this has been shown to help in visual generalization settings (Zhang et al., 2021). We expand on this framework to improve ZSG per- formance, using unsupervised learning to train an agent that recognizes behavior similarity in soward-free fashion. To do so, we propose Cross Trajectory Representation Learning (CTRL), which applies novel self-supervised learning (SSL) objective pairs of trajectories drawn from the agent's s policies. For optimization, CTRL defines a prediction objective across trajectory representations from nearby partitions defined by an online clustering algorithm. The end result is an agent whose encoder maps behaviorally similar trajectories to similar representations without directly referencing reward, which we show improves ZSG performance over using pure RL or RL in conjunction with other unsupervised or SSL methods.
Our main contributions are follows:
* We introduce Cross Trajectory Representation Learning (CTRL), a novel SSL algorithm for RL that defines an auxiliary objective across trajectories in order to capture the notion of behavioral similarity in the representations of the agent's 's belief states. CTRL's approach is two-fold: (i) it uses a clustering loss to group representations of behaviorally similar trajectories and (ii) boosts cross-predictivity between trajectory representations from nearby clusters.
We empirically show that CTRL improves zero-shot generalization in the challenging Procgen benchmark suite (Cobbe et al., 2020). Through series of ablations, we highlight the importance of cross-trajectory views in boosting behavioral similarity.
We connect CTRL to the class of bisimulation methods, and provide sufficient conditions under which both formalisms can be equivalent.
2 BACKGROUND MOTIVATION, AND RELATED WORKS
There are a broad class of ZSG settings in RL, such as generalization across reward functions (Barreto al., 2016; Touati and Ollivier, 2021; Misra et al., 2020), observation spaces (Zhang al., 2021; Li et al., 2021; Raileanu and Fergus, 2021), or task dynamics (Rakelly et al., 2019). For each of these settings, there are number of promising directions for improving ZSG performance: giving the agent better exploration policies (Van Roy and Wen, 2016; Misra et al., 2020; Agarwal et al., 2020), meta learning (Oh et al., 2017: Gupta et al., 2018; Rakelly et al., 2019), or planning (Sohn et al., 2018). In this work, we focus on directly improving the agent's 's representutions. The agent's s representations are high-level abstractions of observations or trajectories from the environment (e.g.. the output of encoder), and the desired property here is that one can casily learn policy top of that representation such that the combined model (i.c., the agent) generalizes to novel situations. The tasks are assumed to share common high-level goal and are set in environments that have the same dynamics, but each task may need to be accomplished under different initial conditions and may differ visually. As the policy is built upon the agent's representations, this motivates the focus of this work for improving generalization: unless the agent's representations generalize well, one cannot expect its policy readily do so.
Unsupervised representation learning has been shown to improve generalization across domains, including zero-shot in vision (Sylvain al., 2020; Wu et al., 2020) and sample-efficincy in RL. (Fy-
2

---

Published as a conference paper at ICLR 2022
senbach et al., 2019; Schwarzer al., 2021; Stooke et al., 2021). In RL, unsupervised objectives can be used as an auxiliary objective (or auxiliary task, Jaderberg et al., 2017). which provide an alternative signal to reward-based learning signal. Due to the potential role the RL loss may play in overfitting (Raileanu and Fergus, 2021), we believe that having a learning objective for agent's representation that is separate from that of its policy is crucial for good 7SG performance.
Self-supervised learning and reinforcement learning. A successful class of models that incorpo- rate unsupervised objectives to improve RL use self-supervised learning (SSL) (Anand et al., 2019; Srinivas et al., 2020; Mazoure et al., 2020: Schwarzer et al., 2021; Stooke et al., 2021; Higgins et al., 2017). SSI. formulates objectives by generating different views of the data, which are essentially transformed versions of the data, generatee by using data augmentation or by sampling patches. While successful in their own way, prior works that combine SSL with RL, do so by applying known SSL algorithms (e.g.. from vision, Hjelm et al., 2018; van den Oord et al., 2019; Bachman et al., 2019; Chen et al., 2020; He et al., 2020; Grill et al., 2020) to RL a nearly off-the-shelf manner, predicting state representations within given trajectory, only potentially using other trajectories as counterexamples in a contrastive loss. As such, these methods' representations can have trouble generalizing latent behavioral patterns present in ostensibly different trajectories.
Bisimulation metrics in reinforcement learning, Our hypothesis is that ZSG is achievable if the agent recognizes behavioral similarity between trajectories based on their long-term evolution. Learning this sort of behavior similarity is a central characteristic of bisimulation metrics (Ferns et al., 2004), which assign value of states which are behaviorally indistinguishable and have the same reward. Reward- -bsed bisimulation metrics have been shown to learn representations that have a num- ber of useful properties, e.g.: smoothness (Gelada al., 2019), visual invariance (Zhang et al., 2021). action equivariance (van der Pol et al., 2020) and multi-task adaptation (Zhang et al., 2020). For 7SG however, encoding relational information based on reward may not actually help (Misra et al., 2020; Touati and Ollivier, 2021; Yang and Nachum, 2021; Agarwal et al., 2021). as the agent may overfit to spurious correlations between high-dimension training. HOMER (Misra et al., 2020) expands on the concept of bisimulation to learn behavioral similarity between states using unsupervised exploration at deployment. Among the existing methods, PSEs (Agarwal et al., 2021) reward-free notion of behavioral similarity is conceptually the closest to CTRL's, and we compare these algorithms empirically in Section 6. However, algorithmically and in terms of their modes of operation, CTRL and PSE are very different. PSE assumes the availability of expert policies for training tasks and learns a representation using trajectories from these experts and action distance measure, which it also assumes to be provided. CTRL doesn't 't make these assumptions and learns a representation online from trajectories simultaneously generated by its substrate RL algorithm.
observations and the reward signal seen during
Mining views across unsupervised clusters. Given our hypothesis that model that learns be- havioral similarities using signal other than reward will perform well on ZSG, there are still many potential models available to learn said similarities in an unsupervised way. A simple and natural choice is to collect agent trajectories as examples of behaviors, then do clustering (online, similar to Asano et al., 2020; Caron et al., 2020) over trajectories. In RL, Proto-RL (Yarats et al., 2021) also uses clustering to obtain a pre-trained set of prototypical states, but for a different purpose - to estimate state visitation entropy in hard exploration problems. However, clustering alone may not be sufficient to recognize behaviors necessary for ZSG, as representations built on clustering only need to partition behaviors, which may bias the model towards similarities evident training experience. This would be counter-produne pronutr-rrdductive for our generalization goal. We therefore use a second objective built on top of the structure provided by clustering to learn a more diverse set of similarities. Drawing inspiration from Mine Your Own View (MYOW, Azabou et al., 2021), CTRL selects (mines) representational nearest neighbors from different, nearby clusters and applies a predictive SSL objective to them. This cross-cluster objective encourages CTRL to recognize a larger set of similarities than would be necessary to cluster on the training set, which we show improves ZSG performance.
3 PROBLEM STATEMENT AND PRELIMINARIES
Formally, we define our problem setting w.r.t. discrete-time Markov decision process (MDP)
M A AS P, R), where S is state space, A is action space, P : S x A S [0, 1] is
3

---

Published as a conference paper at ICLR 2022
transition function characterizing environment dynamics, and R S A R is reward function. M's is state and action spaces may be discrete or continuous, but in the rest of the paper we assume them to be discrete to simplify exposition. In practice, not the full information about the environment's s current state. Consider an observation space O and an observation function 2 S x O [0. 1] that define what observations an agent may receive and how these observations are generated (possibly stochastically) from M's states. We define a task T partially observable MDP (POMIP) 7 (S, A, P. R. O, Z, so), where So E S is an initial state. Although many RL agents make decisions in a POMDP based only on the current observation on or most few recent ones, in general this may require using information from the entire observation history 01,...,0 so far. Denoting the space of such histories as H. computing an agent's 's behavior for task T amounts to finding a policy π H x A - [0. 1] with the optimal or near-optimal expected return from the initial state VF - I. EE= RR,,S, T (Ht) So = so] where St and H; are random variables for the POMDP's underlying state and agent's 's observation history at time step t, respectively, and r is a discount factor. For an MDP M = AS P, R),, a set @ of observation spaces, and set 2 of observation functions w.r.t. S, let a task family be the POMDPs set MM A A, R. O, Z, so) Joee.eE............ in have the same mathematical form, e.g., pixel tensors representing possible camera images, but correspond to qualitatively distinct subspaces of this larger space, such subspaces of images depicting brightly and dimly lit scenes.
agent usually receives observations but
We assume that different observation spaces
Our training and evaluation protocol formalized w.r.t. a tusk distribution di (M....)))))))))) & P(O, 2) over task family FM,,0,2 where P(O, S) is a joint probability mass over observation spaces, observation functions, and initial states. the rest of the paper, M, 0, and will be clear from context, and we will denote the task family as I and the task distribution as d(. For agent training, we choose N tasks from T. and denote the distribution d(.9) restricted to these N tasks as x(Tx). During the training phase, the RL agent learns via series of epochs (themselves composed of episodes), by sampling a task from ((T) independently at the start of each episode, until the total number of time steps exceeds its training budget. In each epoch, an RL algorithm uses a batch of trajectories gathered from episodes in order to compute gradients of an RL objective and update the parameters of the agent's policy In representation learming-aidsd RL, a policy is viewed as a composition π = 0 á of an encoder & H -> E and a policy head 0 E x A - [0, 1]. both of which are the outputs of the training phase. The training phase is followed by an evaluation phase, during which the policy is applied to tasks sampled from d( & \ TNn,, distribution d restricted to the set of tasks 7 \ .N not seen during training. Our focus on generalization means that we seek a policy т whose encoder é allows it to maximize Er-d( gw [V7] despite being trained only on distribution d( x-))
ALGORITHM
CTRL's key conceptual insight is that capturing reward-agnostic behavioral similarity improves ZSG, because it enables é to correctly associate previously unseen observation histories with those for which the agent's 's RL-trained behavior prescribes good action. CTRL runs synchronously with an online RL. algorithm, which is crucial to ensure that as the agent's policies improve, so does the notion of behavioral similarity induced by CTRL.
Like most online RL methods themselves, our algorithm operates in epochs, learning from a batch of trajectories in each epoch. CTRL assumes that all trajectories within each of its training batches come from the same policy. Before the RL algorithm updates the policy head in given epoch, CTRL uses a trajectory batch from the current policy to update the encoder with gradients of novel auxiliary loss LCTRI, that we describe in this section.
4.1 INTUITION AND HIGH-LEVEL DESCRIPTION
Algorithm overview. For each trajectory batch, CTRL performs 4 operations:
1. Apply the observation history (belief state) encoder ø to generate a low-dimensional reward- agnostic representation ("view") of each trajectory.
2. Group trajectories' views into C sets (C is a tunable hyperparameter) using an online clustering algorithm with loss Lehust (Equation (4), Section 4.2).
4

---

Published as a conference paper at ICLR 2022
Clustering
Velast
RL
Concal
112
Prediction
π
V
yood
Opred
W
E
Prediction
Chastering
Anedarr
pred
w
Softmex Ieege
Te:
Sinkhorn
lillllleellz
Mined V.
Gpoed
W
log
Softmar w
P
Leluse
Alneed point
Figure 1: Schematic view of CTRL's key steps for every trajectory batch. (i) Generating trajectory views (top left). For each trajectory in batch, CTRL samples subsequence of its time steps, computes belief-state/actioon embeddings u, with encoder d, and concatenates them into a trajectory representation (view) u. (ii) Clustering trajectory views (bottom right). CTRL uses the online Sinkhorn Sinkhorn-Knopp clustering procedure (Caron et al.. 2020): for each trajectory view it produces two views and soft-clusters all trajectories' vs and into C' clusters, and uses a measure of consistency between these clusterings loss Cahse. the diagram, variables e, denote cluster centroids. (iii) Encouraging cross-cluster behavioral similarity (bottom left). After computing trajectory view clusters, CTRL applies a variant of MYOW (Azabou et al., 2021) to them. Namely, repeatedly samples trajectory view v', computes a new view w' for it, and computes a loss Eyes that penalizes differences between w' and views ve, of randomly chosen trajectories from v's neighboring clusters. Encoder and auxiliary predictors used by CTRL then updated using LCTRL Cetua Lyedd 's gradients (top right).
3. Using trajectory pairs selected from neighboring clusters, apply predictive loss Lpred (Equation (6), Section 4.2) to encourage @ to capture cross-cluster behavioral similarities.
4. Update & with gradients of the total loss: LCTRI Letust + f
The schema in Figure provides a high-level outline of these steps' implementation and explains their interplay within CTRL, accompanied by an intuition for each step (below) more detailed description in Section 4.2. We conduct ablations to show the effect of removing the clustering and predictive objectives of CTRL, with details in Appendix 8.3.
Clustering. C clusters can be viewed as corresponding to C latent "situations" which an RL agent may find itself. Each situation is essentially a group of belief states. CTRL's implicit hypothesis is that a given policy should behave roughly similarly across all belief states corresponding to the same "situation", i.e., generalize across similar belief states. Under this hypothesis, agent's policy can be expected to produce sets of roughly similar trajectories. CTRL's clustering step (#2 above) an attempt to recover these trajectory sets. Since each trajectory consists of belief states, the purpose of Letust is to force the encoder ò to compute belief state representations that make trajectories within each cluster look similar in the latent space.
Since we would like to evolve clustering online as new trajectory batches arrive, we employ a common online clustering algorithm, the Sinkhorn-Knopp procedure (Caron et al., 2020), which has been used as an auxiliary RL loss (e.g., Pro Yarats et al., 2021).
Cross-cluste prediction. Note, however, that the clustering loss emphasizes the recognition of behavioral similarities within clusters. This may hurt gencralization, as the resulting centroids may not faithfully represent behaviors encountered at test time. Our hypothesis is that encouraging encoder 6 to induce latent-space similarities between trajectories from different but adjacent clusters will increase its ability to recognize behaviors in unseen test trajectories.
While there are several ways to encourage cross-cluster representational similarity, using a mechanism similar to MYOW (Azabou et al., 2021) trajectories drawn from neighboring clusters captures this idea particularly well. Namely, to get the cross- cross-predictive loss Lpred, we sample trajectory view pairs from neighboring clusters and apply the cosine-similarity loss to those pairs.
5

---

Published as a conference paper at ICLR 2022
Using reward guidance without reward signal for representation learning. CTRL trains en- coder & only using the gradients of LCTRL: the RL algorithm'ss 's loss CRL. trains only the policy head. Thus, encoder Φ is isolated from the previously observed dangers of overfitting to the reward func- tion (Raileanu and Fergus, 2021) ) that shapes LRL. However, we emphasize that CTRL's representation learning is nonetheless very much guided by the reward function, although indirectly: the training batches of belief state and action trajectories are still collected from policies learned by the policy head via LRL'S 's gradients, which are reward-dependent.
4.2 DETAILS
Below we describe the details of each of CTRL's steps, with CTRL's pscudocode presented in Algorithm 1 in Appendix 8.1. While in general the agent's belief state at step of a trajectory is the entire observation history /4 = (o1..... 0z), in the rest of the section we will assume h = (ot) and, in slight abuse of notation, use ((0:) instead of ((h,) to simplify explanations". We emphasize, however, that CTRL equally applies in settings where the agent uses much longer history of observations as its state. In this case, would be recurrent or process stacks of frames.
We also note that our CTRL implementation's 's high-level algorithmic choices for clustering and cross-cluster prediction Sinkhorn-Knop and MYOW, respectively come from prior works (Caron et al., 2020; Azabou et al., 2021).
Generating low-dimensional trajectory views with encoder o. CTRL's input in each epoch is trajectory batch (rraji/zz of size B. Assume all trajectories in the batch have the same length L. For an integer hyperparameter T < L, for each trajectory trajj- = (oo, do, .o.... OT,, a,, Tr. we independently and uniformly sample a subset of its steps T. = ti....,
.......tz tr to form a subtrajectory (ot,, as,, .t....... Oby atry [tr). We then encode this subtrajectory as
(T) (FiLM ), ) FiLM((((o)))) (t,)).
(1)
where (OLM(((Or,), a.)) is common way of combining representations of different objects, akin to conditioning Perez et al. (2018), and the resulting vector u, (7i) is in low-dimensional space Ur Note two aspects of the process of generating these vectors: (1) drops rewards from the original
trajectory and (2) it critically relies on Φ whose parameter values are learned from previous epochs.
Vectors u: (m) produced in this way are the trajectory views that the next steps of CTRL operate on.
Clustering trajectory views. CTRL groups trajectories from the epoch's batch by clustering the set of their views {u. (x) ) 1= Since we would like to evolve clustering online as new trajectory batches arrive, we employ a common online clustering algorithm, the Sinkhorn-Knop procedure (Caron et al., 2020), which has been used as an auxiliary RL loss (e.g., Proto-RL, Yarats et al., 2021). Since CTRL operates online, Sinkhorn- Knopp is better-suited for the task than other clustering methods.
The clustering branch computes two views of each input u (7:) in a cascading fashion: first by passing it through a clustering encoder t/etusa 4 V, e.g. an RNN, to obtain a lower-dimensional view v; = thetust (T))). and then by passing v; through yet another network, an MIP Aciust V - W, to produce view w. Oetust (v:). The parameters of delust and detsst are learned through the epochs jointly with ø's. Like u!" (T)) each v; and w; is a view of trajectory i; the approach then consists in projecting v.'s 's's and w, 's onto centroids of C' clusters in two different ways and then computes clustering loss that enforces consistency between 8, 's and w. 's cluster projections.
Specifically, represent the centroid of each cluster with vector e V. which are stacked into a matrix E. These vectors are additional parameters in the joint optimization problem CTRL solves. They can be regarded views of C' typical behaviors around which the trajectories' views are regrouped. To project trajectory i's v, views onto behavioral centroids learned from previous trajectory batches, CTRL computes vector of soft assignments of vs to each centroid e.:
e. 1e.l2
Q = Softmax I
|
While, in theory, the Procgen suite is indeed a POMDP, most RL algorithms take the most recent observation as the belief state- a simplification which was shown not to hinder ZSG on Proegen (Cobbe et al., 2020).

---

Published as a conference paper at ICLR 2022
and forms B x C' matrix Q whose i-th row is the soft assignment of vs- The resulting assignments may be very unbalanced, with most probability mass assigned to only few clusters. Applying the Baokhor--Koopp algorithm solves this issue by iteratively re-normalizing Q in order to obtain more equal cluster membership (Cuturi, 2013), where the degree of re-normalization is controlled by temperature parameter (3. The output of this operation is a matrix Q.
For each view w, (the projection of trajectory view 2y) CTRL computes the logarithm of its soft cluster assignments and treats these vectors as rows of another B x C matrix P:
Softmax w eea)
P. log
IsesC
Finally, we compute the entropy between Q and P, which measures their inconsistency. This measure is taken as the clustering loss:
Lelust = CrossEntropy Q P).
(4)
Encouraging cross-cluster behavioral similarity. Note, however, that the clustering loss in the above step emphasizes the recognition of behavioral similarities within clusters. This may hurt generalization, as the resulting centroids may not faithfully represent behaviors encountered at test time. Our hypothesis that encouraging encoder to induce latent-space similarities between trajectories from different but adjacení clusters will increase its ability to recognize behaviors in unseen test trajectories.
While there are several ways to encourage cross-cluster representational similarity, using mechanism similar to MYOW (Azabou et al., 2021) on trajectories drawn from neighboring clusters captures this idea particularly well. Namely, to get the prossctrede loss Cpred, we sample trajectory view pairs from neighboring clusters and apply the cosine-similarity loss to those pairs.
To implement this idea, we define measure of cluster proximity via matrix D of cosine similarities between cluster centroids: for clusters and D& = |leg - ei (Grill et al., 2020). Recall that in the previous step, the clustering branch computed matrix Q whose rows Q; are soft assignments of trajectory i's view v, to clusters. this step, we convert these soft assignments to hard ones by associating a trajectory's 's view v; with cluster 6; = argmax Q; and treating cluster e
argmax, SScc Qi)- To assess how predictive a trajectory embedding u; is of a trajectory embedding uj like in the clustering step will use two special helper maps, d'pred U + V to obtain reduced- dimeedsiinnniitt view o' Upreet ()) and ppred to further project v to w = (p))
consisting of trajectories with indices in the set Te {i
CTRL proceeds by repeatedly sampling trajectories, which call anchor trajectories, from the batch, with their associated embeddings u. For each anchor trajectory n. consider clusters C1...., CK nearest to n's cluster Cu, as defined by the indices of K largest values in row C, of matrix D (we exclude C, itself when determining c, 's nearest clusters). Borrowing ideas from the MYOW approach Azabou et al. (2021), CTRL mines a view for un by randomly choosing trajectory with embedding u from each the neighboring clusters and computing its view v', = v) Plee (u neighbors' views ...,..., v'xx trajectory n''s mined views.
(n)) We call the
For the final operation in this step. CTRL computes trajectory n's predictive view w/, Opeea ( Preed (n)) and measures the distance from it to trajectory n's mined views:
C pred
Iwo v',
(5)
N regulates the number of anchor trajectories to be sampled, so the total prediction loss
N
(6)
pred
n=
Updating encoder e using reward guidance without reward signal for representation learning. Note that CTRL's total loss LCTRI. depends on the parameters of encoder & as well as of clustering
7

---

Published as a conference paper at ICLR 2022
networks Octust and faust prediction networks Pctutt and Betust. and cluster centroids e,, 1.l1sc5 C. In each epoch, CTRL updates all these parameters to minimize LCTRL.
5 CONNECTION TO BISIMULATION
Deep bisimulation metrics are tightly connected to the underlying mechanism of mining behaviorally similar trajectories of CTRL. They operate
space and, as is the case for DeepMDP (Gelada et al., 2019) and DBC (Zhang et al., 2021), ensure that bisimilar states (i.e. behaviorally similar states with identical reward) are located close to each other in that latent space. In this section, we aim to highlight a functional similarity between bisimulation metrics and CTRL.
Definition A bisimilation relation E C S x S is binary relation which satisfies, V(8,)) () E E:
1. Va E a() = R(((
2. Va E A, SANCES,E, P(s,))
a(t,a)(s)
In practice, rewards and transition probabilities rarely match exactly. For this reason, Ferns et al. (2004) proposed smooth alternative to bisimulation relations in the form of bisimulation metries, which can be found by solving recursive equation involving the Wasserstein- distance Wi between transition probabilities. W1 be found by solving the following linear programming (Villani, 2008), where we let r = (v RM 0 e,
S 1 V1 is [V]]
(P(1)) == max WET
(P(ss) ((s))((s) s.t. mss) - mus') < d(s, s')ss E ES,
(7)
#ES
where is a vector whose elements are constrained between 0 and 1. In practice, bisimulation metrics are used to enforce a temporal continuity the latent space by minimization of the W1 loss between training state-action pairs. Therefore, to show connection of CTRL to seward-free bisimulation metrics, it is sufficient to show that two trajectories are mapped to the same partition if their induced Wi distance is arbitrarily small. In (informal) argument that follows, we assume that CTRL samples two consecutive timesteps and encodes them into o; the exact form of v dictates the nature of the behavioral similarity. The proof can be found in Appendix 8.7.
Proposition 1 (Informal) Let M be an MDP where R(s,a) 0 for all (s, a) E S A and let v, o' E V be two dynamics embeddings in M. The clustering operation between u, v induces reward-free bisimilarity metric W1 W1(P()) P(o')) between induced distributions P(u) and IP('").
EMPIRICAL EVALUATION
We compare CTRL against strong RL baselines: DAAC (Raileanu and Fergus, 2021) the current state-of-the-ee on the challenging generalization benchmark suite Procgen (Cobbe et al., 2020), and PPO (Schulman et al., 2017). DAAC optimizes the PPO loss (Schulman et al., 2017) through decoupling the training of the policy and value functions, which updates the advantage function during the policy network updates. We then compare to several unsupervised and SSL auxiliary objectives used in conjunction with PPO. DIAYN (Eysenbach et al., 2019) is an unsupervised skill- based exploration method which we adapt to the online setting by uniformly sampling skills. Its notion of skills has some similarities to the notion of clusters in CTRL. We also compare with two SSL-based auxiliary objectives: CURL (Srinivas et al., 2020), common SSL baseline which contrasts augmented instances of the same state, and Proto-RL RL (Yarats et al., 2021), which adapt for this generalization setting. Finally, we provide a comparison against bisimulation based algorithms: DBC (Zhang et al., 2021), which was shown to perform well on robotic control tasks with visual distractor features, and PSE (Agarwal et al., 2021). PSE assumes policies for training tasks to be given and, like Agarwal al. (2021), we it both with random and high-quality policies pretrained with extra computation budget. See Appendix 8.2 for details.
in our experiments, consists of 16 video games (see Table 1). Procgen procedurally generates distinct levels for each game. The number of levels for each game is virtually unlimited. Levels within game which share common game rules and objectives
The Procgen benchmark suite, which
8

---

Published as a conference paper at ICLR 2022
but differ in level design such the number of projectiles, background colors, item placements throughout the level and other game assets. All of this makes Procgen a suitable benchmark for zero-shot generalization. Using
notation from Section 3, for each of 16 games, we train uniform distribution d( ((Tx) over = 200 "easy" levels of the game and evaluate on d( T Ix), i.e., uniform distribution over the game's "easy" levels not seen during training. Following Mohanty et al. (2021), we report results after 8M steps of training, since this demonstrates the quality of ZSG that various representation learning methods achieve quickly. However, Figure 4 in Appendix 8.3 also provides results after 25M steps of training, as in the original Procgen paper (Cobbe el al., 2020).
RL
RLaBisin
RLaUnsup.
RLASSL
Ours
Env
PPO
DAAC
PPOxDBC
PPO+PSE (random)
PPO:PSE (pretrained)
PROYDIAYN
Proto-RI.
PPO,CURL
CTRL
higfish
bossdight
5.2:0.33
1.720.7
$20.1
0.2
0,7 0.1
1.1:0.2
6.1210.5
4.620..8
8.2:0.11
caveflyer
4.420.3
4.3.0.11
1.610.1
1610.1
36 1611.3
1.011.9
4,711..1
4.6.1..3
4.710.22
choser
7.2::022
7.110.1
1.800.1
8,7 1
0.2
5.6.10.5
7.6.0.22
7.2:10.2
1.1.0.22
climber
5.1:0.11
5.5:10.2
411104
192
0.8+0.7
5.5:0.3
5.5:0.1
5.9:11.2
coiorun
79+001
81++0.1
dodgehall
1.3=0.1
18:+0.2
1.0+0.3
16=0.1
13 0.1
14#002
L6++0.1
1.4+0.1
1.++0.1
fruihot
12.4:00.2
11.5:10.3
7.62002
1.
11 0.2
1.243.0
12.3=0.4
12:3=00.2
13.3:0.3
beist
2.72:0.2
3.4.0.2
3.31022
3.1 1.3
1.4
0.210.2
3.010.3
2.510.1
3.1
jumper
5.810.3
6.3.0.1
53102
5.9.0.11
6,010.11
loger
3.5:004
3.5:004
2.7:0.1
27 07
2610.1
2,520.2
3,2008
3.6:0.5
2.8:0.2
maze
5.4-0.2
5.6+0.22
50+00.1
54+022
53+001
1.6+122
5.5+0.3
5.4+0.11
5.7+1.1
miner
8.7-033
5.7+0.1
48+0.1
5.=
43 43
1,3+200
8.8+0.5
8.6-0.2
6,5+0.2
5.5:0.2
5.2:0.11
3.520.1
34
0.
2.82200
5,42033
5.620.1
58:0011
plunder
starpilot
4.72002
4.11022
2810..1
32102
3.0
5.810.66
5,210.2
5.020.1
7.7.0.55
Table 1: Average evaluation returns collected after 8M training frames, one standard deviation over 10 seeds.
Main results. As Table I shows, PPO+CTRL outperforms all other baselines, including DAAC, on most games. Notably, disimulation-basee approaches other than CTRL- DBC as well PSE with both random and expert data data-gathering policies exhibit lower gains than others. While this can be surprising, recent work has seen similar results when applying DBC to tasks with unseen backgrounds (Li et al., 2021). PSE's inferior performance may be due to the policy similarity metric, which PSE requires as imput and which we took from Agarwal al. (2021), being poorly suited to Procgen. This highlights an important difference between CTRL and PSE: CTRL doesn' 't need policy similarity metric, since implicitly induces such a metric based on trajectory "signatures".
Despite training static prototypes for 8M timesteps and adapting the RL head for 8M additional timesteps (see Appendix 8.2 for details), Proto-RI. performs worse than CTRL. This suggests that the temporal aspect of clustering is key for ZSG, hypothesis we explore further in Section 8.4. Likewise, PPO+DIAYN uses its pre-training phase to find a diverse set of skills, which can be useful in robotics domains, but does not help much in the ZSG setting of Procgen. DAAC also exhibits good generalization performance, but inherits from PPG (Cobbe et al., 2021) the separation the value and policy branches both parameter-wis wiise and by introducing distinct training phases, overhead which CTRL manages to avoid. In addition, we describe a number of ablation studies (Appendix 8.3), empirically show that slow clustering convergence leads to better generalization (Appendix 8.4), and demonstrate on toy example how learning behavioral similarities captures local perceptual changes (Appendix 8.5).
CONCLUSIONS
This work proposed CTRL, novel representation learning algorithm that facilitates zero-shot generalization of RL policies in high-dimensional observation spaces. CTRL can be viewed as inducing unsupervised reward-agnostic bisimulation metric over observation histories, learned over transitions encountered by policies from an RL algorithm's value improvement path (Dabney al., 2021). We hope that in the future CTRL will inspire other representation learning methods based on capturing belief states's 's behavioral similarity, which will be capable of policy generalization across greater variations in environment dynamics.

---

Published as a conference paper at ICLR 2022
REFERENCES
A. Agarwal, M. Henaff. S. Kakade. and W. Sun. Pc-pg: Policy cover directed exploration for provable policy gradient learning. In NeurIPS, 2020.
R. Agarwal, M. C. Machado, embeddings for generalization in reinforcement learning. In ICIR, 2021.
Castro, and M. Bellemare. Contrastive behavioral similarity
A. Anand, E. Racah, S. Ozair, Y. Bengio, M.-A. Côté, and R. D. Hjelm. Unsupervised state representation learning in Atari. In NeurlPS, 2019.
Y. M. Asano, C. Rupprecht, and A. Vedaldi. Self-labelling via simultaneous clustering and represen- tation learning. ICLR, 2020.
M. Azabou, M. G. Azar, R. Liu, C.-H. Lin, C. Johnson, K. Bhaskaran- M. Dabagia, Hengen, W. Gray-Roncal, M. Valko, and E. L. Dyer. Mine your view: Self-supervised learning through across-sample prediction. arXiv preprint arXiv:2102.101100 2021.
Bachman, R. D. Hjelm, and W. Buchwalter. Learning representations by maximizing mutual information across views. In NeurlPS, 2019.
A. Barreto, W. Dabney, Munos, Hunt, T. Schaul, H. Van Hasselt, and D. Silver. Successor features for transfer in reinforcement learning. arXiv preprint arXiv:1606.05312. 2016.
M. Bucher, S. Herbin, and F. Jurie. Generating visual representations for zero-shot classification. In Proceedings of IEEE ICCV Workshops, 2017.
M. Caron, I. Misra, J. Mairal, P. Goyal, P. Bojanowski, and A. Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In NeurIPS, 2020.
T. Chen, Kornblith, M. Norouzi, and G. Hinton. A simple framework for contrastive learning of visual representations. In ICML, 2020.
K. Cobbe, C. Hesse, J. Hilton, and J. Schulman. Leveraging procedural generation to benchmark reinforcement learning. In ICML, 2020.
K. Cobbe, J. Hilton, O. Klimov, and Schulman. Phasic policy gradient. In ICML, 2021.
M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transportation distances. In NeurIPS, 2013.
W. Dabney, Barreto, M. Rowland, R. Dadashi, Quan, M. Bellemare, and D. Silver. The value-improvement path: Towards better representations for reinforcement learning. In AAAI, 2021.
B. Eysenbach, A. Gupta, J. Ibarz, and S. Levine. Diversity is all you need: Learning skills without a reward function. In ICLR, 2019.
N. Ferns, P. Panangaden, and D. Precup. Metrics for finite Markov decision processes. In UAI, 2004.
C. Gelada, Kumar, Buckman, O. Nachum, and M. Bellemare. DeepMDP: Learning continuous latent space models for representation learning. In ICML, 2019.
J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. H. Richemond, E. Buchatskaya, C. Doersch, B. A. Pires, Z. D. Guo, M. G. Azar, B. Piot, K. Kavukcuoglu, R. Munos, and M. Valko. Bootstrap your own latent: A new approach self-supervised learning. In NeurIPS, 2020.
A. Gupta, B. Eysenbach, C. Finn, and S. Levine. Unsupervised meta-learnin for reinforcement learning, arXiv preprint arXiv: 04640, 2018.
K. He, H. Fan, Y. Wu, S. Xie, and R. Girshick. Momentum contrast for unsupervised visual representation learning. CVPR, 2020.
I. Higgins, A. Pal, A. A. Rusu, L. Matthey, C. P. Burgess, A. Pritzel, M. Botvinick, C. Blundell, and A. Lerchner. DARLA: Improving zero-shot transfer in reinforcement learning. In ICMI, 2017.
10

---

Published as a conference paper at ICLR 2022
R. D. Hjelm, A. Fedorov, S. Learning deep representations by mutual information estimation and maximization. In ICLR, 2018.
K. Grewal, P. Bachman, A. Trischler, and Y. Bengio.
M. Jaderberg, V. Mnih, W. M. Czarnecki, T. Schaul, J. Z. Leibo, D. Silver, and K. Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. ICLR, 2017.
A. Kendall, J. Hawke, D. Janz, P. Mazur, D. Reda, J.-M. Allen, V.-D. Lam, A. Bewley, and A. Shah. Learning to drive in a day. In ICRA, 2019.
S. Levine, C. Finn, T. Darrell, and P. Abbeel. End-to-end training of deep visuomotor policies. Journal of Machine Learning Research, 17(39):1-40,, 2016.
B. Li, V. François-Lavet, T. Doan, and Pineau. Domain adversarial reinforcement learning. arXiv preprint arXiv:2102.0070 2021.
B. Mazoure, R. T. des Combes, T. Doan, P. Bachman, and R. D. Hjelm. Deep reinforcement and infomax learning. In NeurlPS, 2020.
D. Misra, M. Henaff, A. Krishnamurthy, and J. Langford. Kinematic state abstraction and provably efficient rich-observation reinforcement learning. In ICML, 2020.
V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Ried- miller, A. K. Fidjeland, G. Ostrovski, S. Petersen, C. Beattie, A. Sadik, I. Antonoglou, H. King, D. Kumaran, Wierstra, Legg, and D. Hassabis. Human-level control through deep reinforce- ment learning. Nature, 518(7540):529- 533,3 2015.
S. Mohanty, J. Poonganam, A. Gaidon, A. Kolobov, B. Wulfe, D. Chakraborty, G. Semetulskis, Schapke, J. Kubilius, Pašukonis, al. Measuring sample efficiency and generalization in reinforcement learning benchmarks: NeurlPS 2020 Procgen benchmark. arXiv preprini arXiv:2103.15333 2021.
J. Oh, S. Singh, H. Lee, and P. Kohli. Zero-shot task generalization with multi-task deep reinforcement learning. ICML, 2017.
E. Perez, Strub, de Vries, V. Dumoulin, and A. Courville. Film: Visual reasoning with general conditioning layer. In AAAI, 2018.
R. Raileanu and R. Fergus. Decoupling value and policy for generalization in reinforcement learning. In ICML, 2021.
K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen. Efficient off- off-policy meta-reiiforcemment learning via probabilistic context variables. In JCML, 2019.
J. Rousseeuw. Silhouettes: a graphical aid the interpretation and validation of cluster analysis. Journal of computational and applied mathematics, 20:53-65, 1987.
J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint /rXiv: 1707.0637, 2017.
M. Schwarzer, A. Anand, R. Goel, R. D. Hjelm, A. Courville, and P. Bachman. Data-efficient reinforcement learning with self-predictive representations. ICLR, 2021.
S. Sohn, J. Oh, and H. Lee. Hierarchical reinforcement learning for zero-shot generalization with subtask dependencies. arXiv preprint arXiv:1807.07663 2018
A. Srinivas, M. Laskin, and P. Abbeel. CURL: Contrastive unsupervised representations for rein- forcement learning. /CML,, 2020.
A. Stooke, K. Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforcement learning. In ICML, 2021.
T. Sylvain, L. Petrini, and D. Hjelm. Locality and compositionality in zero-shot learning. In ICLR, 2020.
11

---

Published as a conference paper at ICLR 2022
A. Touati and Y. Ollivier. Learning representation optimize all rewards. In ICLR Self- supervision for Reinforcement Learning Workshop. 2021.
A. van den Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprini arXiv: arXiv:8877.00 007.00348, 2019.
E. van der Pol, T. Kipf. A. Olichock, and M. Welling. Plannable approximations mdp homomor- phisms: Equivariance under actions. In AAMAS, 2020.
B. Van Roy and Z. Wen. Generalization and exploration via randomized value functions. In ICML, 2016.
C. Villani. Optimal transport: old and new. Springer Science & Business Media, 2008.
J. Wu, T. Zhang, Z.-J. Zha, J. Luo, Y. Zhang, and F. Wu. Self-supervised domain-awan generative network for generalized zero-shot learning. In CVPR, 2020.
M. Yang and O. Nachum. Representation matters: Offline pretraining for sequential decision making. In ICLR, 2021.
D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Reinforcement learning with prototypical representa- tions. In ICMI., 2021.
A. Zhang, S. Sodhani, K. Khetarpal, and J. Pineau. Learning robust state abstractions for hidden- parameter block mdps. In ICLR, 2020.
A. Zhang, R. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant representations for reinforcement learning without reconstruction. In ICLR, 2021.
12

---

Published as a conference paper at ICLR 2022
8 APPENDIX
CTRL PSEUDOCODE
Algorithm I: Cross Trajectory Representation Learning
Inputs
:online encoder ø, cluster projector Oclust. cluster encoder t'elust. mining projector pred mining encoder prred cluster basis matrix E Hyperparameterss B trajectory batch size, C num. of trajectory clusters, T subtrajectory length, K num. nearest clusters for view mining, - trajectory length, N num. of anchors for view mining, 3 Sinkhorn temperature
for each iteration itr = 1.2... do for each minibatch B do for each trajectory Ti in B do
11, - IT ~Uniform(T)) // Sample temporal keypoints (T) (FiLM( )
cluster dynamics
U, = [u(ti), u [T))) // batch dynamics
U; Velust u;) // fetch embeddings
W. = Betust U;) // fetch projections
U 10. // normalize embeddings
Q Softmax (v, E/3) // Compute latent dynamics scores
10
11
Q Sinkhorn Q) // no cormaiize scores through Sinkhorn
log [Softmax E 8)] // Compute projected dynamics scores Celust (ó, Veluster, feluster) CrossEntropy(Q..P.
12
13
// Predicting neighbors
14
Do lle e,1 // find pairwise basis distances
15
Pped = 0
16
for each anchor = 1,.., do
Tj B // Sample anchor trajectory
U, u/5) // Set anchor embedding
(1)
(k) = ton--kn k.) /1 Find nearby clusters
... C
Ua Uag ~ p(u, P(.,, // Sample views from clusters
29
21
0 ..... Veex Uppeed (Me,) mm ppd (ue,) // embed mined views
23
w/ = Opreaa (eeree (u.)) // min mining target
c(m) Ek=1 1|', StopGrad(e )
pred
Sped = N 1==1= ((n))
25
26
LCTRI- = Lelust + Lped // update networks
Φ, Vesuust Octusts Opeed; Upred Adam(ó; veelust Oelusty Opreds Upedd LCTRI.)
27
for each minibatch B do
29
n, V Adam n. V*; (B)( // update RL parameters
8.2 EXPERIMENT DETAILS
We implemented all algorithms on top of the IMPALA architecture, which was shown to perform well on Procgen (Cobbe et al., 2020). the Procgen experiments, Proto-RL was ran without any intrinsic rewards (since the domains and RL losses for 8M timesteps, after which only the RL loss was optimized for an additional 8M steps (16M steps in total). Similarly, DIAYN also with the pre-training phase of 8M and then RL only objective for the second 8M phase.
not exploration-focused) by first jointly training the representation
Like Proto-RL and DIAYN, PSE needed extra training budget and extra adjustments for fair comparison to CTRL (see Section 2). Agarwal et al. (2021) used PSE only with Soft Actor-Critic, which doesn't it perform well Procgen. Therefore, carried over PSE' available implementation
13

---

Published as a conference paper at ICLR 2022
Name
Description
Value
X
Discount factor
0.999
Decay
0.95
Reimesteps
Number of timesteps per rollout
256
eepochs
Number of epochs for RL and representation learning
1
nsaaples
Number of samples per epoch
8192
Entropy bonus
Entropy loss coefficient
0.01
Clip range
Clip range for PPO
0.2
Learning rate
Learning rate for RL and representation learning Number of parallel environments
5 x 10
Number of environments
32
Optimizer
Optimizer for RL and representation learning
Adam
Frame stack
Frame stack X Procgen frames
1
E
Number of clusters
200
k
Number of k-NN nearest neighbors
3
T
Number of clustering timesteps Clustering temperature
2
B
0.3
Table 2: Experiments' parameters
codebase, with the help of PSE' authors, combine it with the same PPO implementation that CTRL used.
from https://agarwl.giin wtt.:a/www..gtuu..//sseeeee
PSE assumes heing given policies for training problem instances (Procgen levels). Like Agarwal et al. (2021), we ran PSE using both random and high-quality pretrained policies for these problems. In the latter case, we pretrained expert policies for the first 40 levels of Procgen to generate training trajectories for PSE. Each level's expert was trained on 0.5M environment steps. Note that this is much less than 8M steps used for policy training in experiments with other algorithms, but this is because each expert needed to be good only for single level, and we verified that they indeed were. We did this only for the first 40 levels because even for 40 levels this took 20M training steps per game and had to be done for 16 games. We don't 't believe more than 40 experts per game would have made a difference.
Given policies for the training levels, PPO+PSE's training on Procgen mimicked CTRL's: PPO+PSE trained by interacting with the first 200 levels for 8M steps and was evaluated on the rest. However, during the training, PPO+PSE sampled an additional 1M interactions from the pretrained traiing-level policies. the process was repeated for all 16 games, for 10 seeds each.
We did hyperparameter grid search on PPO+PSE's 's hyperparameters for PSE - loss coefficient values of (0.1, 1, 2) and temperature (0.1, 0.3, 0.7). PPO's s hyperparameters were the same as CTRL.
Thus, due to the need to pretrain and gather data with per-level expert policies, PPO+PSE received 0.5 20M 1M M extra environment interactions compared to CTRL, i.e., used 21 21M//M 2.6x more training data than the latter.
8.3 ADDITIONAL RESULTS
ZSG over 8M and 25M training steps and policy performance throughout training We provide the performance plots of various representation learning and RL algorithms for the 8M and the 25M benchmarks, which test zero-shot generalization under different sample regimes. Figure shows the training performance of agents 8M frames, Figure 3 the test performance of agents 8M frames and all levels, and, finally, Figure 4 shows the test performance agents on 25M frames and all levels.
Ablations on algorithm components We ran multiple versions the algorithm to identify the key components which make CTRL perform well in Procgen. The first modification, CTRL consecutive consists in rumning our algorithm but sampling consecutive timesteps, that is f, = t1 + for all t and 0 < T. The second modification, CTRL action removes the action conditioning layer in the log-softmax probability p. and the cluster scores to test the importance of action information for cluster membership prediction. The third modification, CTRL. no cluster removes the clustering
14

---

Published as a conference paper at ICLR 2022
bigtish
bassfight
14
12-
10
10
6
4
4-
2-
dodigeball
10-
8-
15
10
7
6
2
heist
junt per
leaper
maze
4-
Acorage
3
0-
miner
ninja
p plunder
starpiltt
12.5
10
10.0
5.0
4-
0.0
2-
2
5
1 3
7
DAAC
Ours
PPO
PPO+CURL
PPO+
PPO+Sinkhom
Figure 2: Training results over the 8M frames benchmark.
loss, and only restricts to mining and predicting nearby neighbors in the batch. Finally, the last modification, CTRL no pred removes the loss predicting samples from neighboring partitions and only relies the clustering loss to update its representation.
Results suggest that (1) using consecutive timesteps as for the dynamics vector embedding yields lower average rewards than non-consecutive timesteps, (2) action conditioning helps the agent to pick up on the local dynamics present in the MDP and (3) both clustering and predictive objectives are essential the good performance of our algorithm. Results of the last column computed over 10 seeds, rest over seeds.
Ablation number of clusters and clustering timesteps How should one determine the optimal number of clusters in a complex domain? Can the number of clusters be chosen a priori running any training?
Below, we provide some partial answers to these questions. First, the optimal (or truc) number of clusters is domain-specifie, Second, the length of the clusters, i.e. the number of trajectory timesteps passed to Sinkhorn-Knnn widely impact the nature of learned representations, and hence the downstream performance of the agent.
it depends the exact connectivity structure of the MDP at hand.
Temporal connectivity of the clusters Are clusters consistent in time for given trajectory? To verify this, we trained CTRL on million frames of the bigfish game. For every states given trajectory, we have computed the hard cluster assignment the nearest cluster, which yields
15

---

Published as a conference paper at ICLR 2022
bigtish
bassfight
10-
10
8-
6-
2
dodgeeblll
8
2.2
15
8-
1.8
1.5
2
5-
heist
junt per
leeper
maze
6
Arorgge
2-
miner
ninja
p plunder
starpilat
10
6-
5
6
3
0
5
1
4 7
DAAC
Ours
PPO
PPO+CURL
PPO+
PPO+Sinkhom
Figure 3: Evaluation results over the 8M frames benchmark.
Table 3: Average evaluation returns collected after 8M of training frames, + one standard deviation.
Env
CTRL consecutive
CTRL no action
CTRL no cluster
CTRL no pred
CTRL
bigfish
3.9+0..3
3.2+0.33
2.5+0.33 7.8+0.33
3.7+0.1
4.7+0.2
bossfight
8.9+0.1
6.9+0.9
6.6+0.88
8.2+0.1
caveflyer
4.640..2
4.7.0.11
4.640.1
4.641..1
4.710.22
chaser
7.410.3
6.710.2
7.0+0.55
6.510.11
7.110.2
climber
6,210.4
5,5:0.1
5.310.4
5.710.44
5,910.2
coinrun
8.8+0.1
8.5+0.3
8.1+0.2
8.440.3
8.7+0.3
dodgeball
1.8+0.1
1.7+0.1
1.7+0.2
1.7+0.1
1.8+0.1
fruithot
13.1+0.33
12.9+0.4
13.0+0.55
12.5+0.6
13.3+0.3
heist
3.010.1
3.210.3
3.2:0.1
3.010.2
3.1
jumper
6.110.2
6.010.11
5.910.1
5,910.1
6.0.1..1
leaper
3.41.11
3.240.4
2.6+0.3
3.3+0.2
2.8+0.2
maze
5.6+0.2
5.6+0.1
5.740.1
5.8+0.1
5.7+0.1
miner
7.0+0.9
5.940.4
5.6+0.1
6.0+0.2
6.5+0.2
ninja
5.710.1
5.310.2
5.510.1
5.610..1
5.810.11
plunder
6.410.2
5.610.11
5.910.5
6.110.2
6.6.1.33
starpilot
7.0.0.22
4.910..6
4.911.22
5.810.4
7.710.5
a sequence of partitions. We then computed the cosine similarity between time-adjacent cluster centroids, the metric reported the smoothed graph below.
16

---

Published as a conference paper at ICLR 2022
bigfish
bossfight
caveflyer
chaser
Test returss
1.5
5.0
0.0
climber
coinrun
dodgeball
241
fruitbot
Test returns
6-
2
0-
heist
jumper
leaper
maze
1.0
6-
1.0
0.8
0.8
Test
1.00
ninja
plunder
2e1
starpii
0.50
Test
0.25
2
10 Timesteps (M)
Timesteps (M)
Timesteps (M)
Timesteps IM)
CTHLL
Figure 4: Evaluation results over the 25M frames benchmark.
bigfish
leaper
starpilot
Average test reward
3.4
3.5
5
3.2
3.0
3.0
4
2.5
2.8
VW mmm
2.0
2.6
3
1.5
2.4
0
1
4
5
7
0
3
4
5
7
0
1
4
5
Timesteps (M)
Timesteps (M)
Timesteps (M)
Clustering timesteps
Figure 5: Ablation on the clustering timesteps used in the dynamics embedding
Average age test reward
bigfish
leaper
starpilot
4.5
4.5
4.0
3.5
3.0
2.5
2.0
1.5
1
5
6
7
Timesteps (M)
Timesteps (M)
Timesteps (M)
Number of clusters
100
300
Figure 6: Ablation the number of clusters used in CTRL.
Loss landscape of Letsst and Lpred Works relying non-colinear signals, e.g. behavioral similarity and rewards, as is the case for DeepMDP (Gelada et al., 2019), show that interference can occur between various loss components. For example, (Gelada al., 2019) showed how their dynamics and
17

---

Published as a conference paper at ICLR 2022
bigfish
mwm leaper
starpilot
Cluster similarity
12.5
20
10.0
30
15
7.5
20
5.0
10
10
2.5
5
0.0
0-
2.5
Timesteps (M)
Timesteps (M)
Timesteps (M)
Figure 7: Average trthin-trajectory cluster similarity over 1M consecutive timesteps.
bigfish
bigfish
5.0
1.50
1.25
4.8
Celust
pred
1.00
0.75
4.4
0.50
0.25
8
Timesteps (M)
Timesteps (M)
Figure 8: Average values of Calust and Lered over time.
reward losses are inversely proportional to each other early on in the training, taking considerable amount of frames converge.
We observe similar pattern in Figure 8: the clustering loss first jumps while the predictive loss is minimized, then the trend reverses, and both losses get minimized near the end of the training.
Case study: splitting other dynamics-aware losses Similar to the postulate ATC (Stooke al., 2021), we hypothesize that training the encoder only with the representation loss has the most beneficial effect when the representation loss contains information about dynamics. To validate this, conducted an additional set of experiments which leverage predictive information about future timesteps: Deep Reinforcement and InfoMax Learning (DRIML) (Mazoure al., 2020) and Self-Predictive Representations (SPR) (Schwarzer et al., 2021). l). We ran (i) the default version of the algorithms with joint RL and representation updates, well (ii) RI. updates propagated only through the layers above the encoder.
well-known self-supervised learning algorithms
Qualitative assessment of clusters Figure shows, for environments, 4 randomly sampled states for 2 behavioral clusters (4 clusters for Starpilot differences between clusters are easier to visualize in this environment). Note that clustered states go beyond visual similarity, and capture action sequences. agent position, presence of enemies and even topological equivalence of various levels. The choice of environments for this demonstration is dictated by the nature of the action space, c.g. projectiles in Starpilot and path tracing in Miner allow better visualize agent's 's behavior. Note that, for Bigfish, CTRL implicitly picks up the notion of reward density by learning to separate states abundant of fish from those without fish (due to the policy being trained on rewards and thus exhibiting different behavior in those two settings).
Figure 10 shows the t-SNE of from randomly sampled states along the CTRL training path Starpilot - embeddings learned by CTRL. can be seen to concentrate into distinct clusters and around their respective centroids.
can
18

---

Published as a conference paper at ICLR 2022
Env
DRIML
SPR
bigfish
+0.17
+0.15
bossfight
+0.56
+10.36
caveflyer
+0.23
-0.02
chaser
-0.14
-0.13
climber
0.15
+0.14
coinrun
-0.37
+2.05
dodgeball
-0,39
+0.12
fruitbot
-0.53
-0.1
heist
-0.3
-0.03
jumper
+0.04
+0
leaper
-0.04
+0.08
maze
+0.01
-0.15
miner
+0.88
+0.03
ninja
-0.13
-0.31
plunder
-0.02
-0.18
starpilot
-0.11
-0.15
Norm. score
+0.01
+0.74
Table Normalized improvement scores split updates over joint updates encoder, averaged over random seeds.
Starpilot move mant
Miner lesploratory
Bossfight
Jelnecied
Clurder lowrying movement spree))
Cluster Nocal bahavior
Custer trotation firs)
Bigfish
Climber platforms on right sids)
(coordinated five pattems, screant
(yare spanee rewvards
Cluster & mmid-screnn sporadic fie pattems)
Cluster 2 labundand fish dense revord)
Cluster lscarse platforms)
Figure 9: Sample states from behavioral clusters found by CTRL after 2M of training frames for 5 representative environments. The two gray squares in top left is added to indicate the agent's velocity.
8.4 SHOWCASE: SLOW CLUSTERING CONVERGENCE LEADS TO BETTER GENERALIZATION
Trajectory clustering key to representation learning not only CTRL but also prior method, Proto-RI. However, while Proto-RI. uses it to pretrain representation which it then keeps frozen during RL, CTRL applies clustering to evolve the representation as RL progresses. This raises a question: how important is online clustering convergence rate for learning good representation? Intuitively, online clustering converges too quickly and behavioral similarities are "pinned down" early in the training process, the resulting representation will not be robust to distribution shifts induced by improved policies. Therefore, it seems crucial to learn the behavioral similarities at a rate that allows cluster centroids to adapt to the value improvement path (Dabney al., 2021). To validate this hypothesis, we conducted an analysis of the correlation between the clustering quality and the test performance function of training progress three representative games.
19

---

Published as a conference paper at ICLR 2022
0.5M
1.0M
60
1.3M
2.0M
40
20
0-
-0
Figure 10: t-SNE of 1024 randomly sampled states from data collected CTRI. after 0.5,1,1.5 and 2M frames Starpilot, with 0.1, 7 4. As learning progresses, agent behavior clusters become more and more distinct.
To measure the clustering quality, we report the silhouette score (Rousseeuw, 1987), a commonly used unsupervised goodness-of-fit measure which balances inter- and intra-cluster variance. Results shown
bigfish
leaper
starpilot
Silhouette score
0.8
0.6
0.4
0.4
0.2
0.2
0.0
Average test reward
bigfish
leaper
starpilot
4.5
4.0
3.5
3.0
2.5
2
3
4
5
6
2
3
4
5
7
1 2
Timesteps (M)
Timesteps (M)
Timesteps (M)
8-0.01
$-0.1
p=0.2
pm0.3
Figure 11: Goodness-of-clusteree of-lusstriggn measured by silhouctte scores (top) and average test returns (bottom) as function of training samples.
in Figure 11 provide evidence that picking the unsupervised learning procedure which converges the fastest (i.c. uses the lowest temperature () does not necessarily lead to the best generalization performance. Based on Figure 11, we conjecture that fast clustering convergence hinders the performance of the RI. agent due to clusters being fixed early and not adapting to the distribution shift induced by the evolving RL policy.
8.5 SHOWCASE: LEARNING BEHAVIORAL SIMILARITIES CAPTURES LOCAL PERCEPTUAL CHANGES
To demonstrate the importance of identifying behavioral similarities, designed toy example problem with behavioral clusters, where clustering the behaviors correctly leads to finding a near-optimal policy.
20

---

Published as a conference paper at ICLR 2022
Our example problem is based on the standard Ising model* a 32 32 binary lattice, each entry of which evolves at every timestep according to the values of its neighbors, with strengths of neighbor dependencies being regulated by temperature parameter 1/8. We randomly initialize 5 Ising models, each parametrized by an inverse temperature parameter on a uniform grid /3 E [0.01, 0.3). The system state is given by the state of all 5 models, and all models evolve in parallel at every step. At each timestep, the agent needs to choose one of the models, and has 5 actions corresponding to these choices. At timestep i, the agent is allowed to observe only the state of the model it chose this step, and gets reward based only on this model's state. The reward yielded by Ising model at timestep t given by r, - -Issi, G where G' is goal state. For given problem instance, sample G randomly by instantiating a 6th Ising model with an unknown inverse temperature parameter /3' and letting G be its final configuration after evolving it for a random number of steps. Figure 12 outlines the experimental setting for this study case.
The behavioral clusters in our setting correspond to the Ising models. The optimal strategy to solve this task is to 1) identify the Ising model (i.e., the behavioral cluster) whose temperature parameter is closest to 3 and 2) choose that model and collect the corresponding rewards.
Branch
Branch
Goal configuration
Agent
Action
B 3-2
8 = 2.5
Ttt = 8t
Figure 12: The composite Ising matching problem: agent has match a given Ising configuration by swapping branches various transition dynamics
Table outlines the results obtained by deploying CTRL with different number-of-clusters parameter value. One can see that the largest improvement in silhouette score occurs from E = to E = 5 (14.5%), suggesting that monitoring the largest change in silhouette can be used to set the true number of clusters in CTRL which, in turn, corresponds to the highest-return policy discovered by CTRL.
= 2
E
=
= 6
= 50
Returns Silhouette
-0.78
-0.96
-0.02
-0.07
-0.54
0.875
0.796
0.651
0.554
0.039
Silhouette change
0.079
0.145
0.097
0.515
Table 5: Returns and silhouette scores obtained by CTRL in the composite Ising matching domain.
8.6 ADDITIONAL THEORETICAL FINDINGS
Do uncorrelated local changes to state embeddings affect the clustering?
Theorem 8.1 Let M be an MDP and let E V be dynamics embeddings in M. Define
5, = 5 1 s |
(8)
5 h<isTW TIV
and pick 8 s.I. it lies on the positive half-plane spanned by E. EJ for some S S E. Then, o and v belong to the same partition j.
It becomes apparent from the above statement that perturbations to single state groups of state embeddings do not modify the partition membership as long as their direction aligns with that the cluster embeddings.
21

---

Published as a conference paper at ICLR 2022
8.7 PROOFS
Throughout this section, we assume that the policy x is fixed, and that CTRL optimizes Cetust only.
Proof (Theorem 8.1) For two dynamics embeddings to be assigned to the same cluster j. the following should hold
M Σ
MV
J
2
j':
1m
(9)
Us + 6 E
(vs + 6:) 6
for any S S s.s. j.
M
M
V
M4
E J
E
+
6,E j/
2=1
i=1
i=1
(10)
[
E
E
(E
E
Taking the difference between both equations yields the necessary condition for two dynamics to belong to the same cluster
sup ISJ'SE
E
B 6; 2 0, S <
(11)
Corollary 8.1.1 Let v, u' be two dynamics embeddings, and define ô v If belongs to cluster j and = arg m Ejő, then v' also belongs to cluster j.
Perturbations are of the form E E.J )6 1 then the cluster assignment doesn't change. Let be of size kh [V.. Define, without loss of generality
6: 5 5 1 isish h
(12)
6, = 8
h<ish kh
and pick 8 s.t. it lies on the positive half-plane spanned by E.J E.yr- Then,
M
E.
>(E))
[E.
E. E.y0) )o > (D., E.y )0 Σ (E.
E.
2 (13)
which concludes the proof.
Proof 2 (Theorem 1) Since the Wi metric is defined between distribution functions, throughout the proof to denote the probability distribution over elements of the dynamics vector v. In practice this amounts to re-normalizing the representation.
P|u]
For two dynamics to be assigned the same cluster j, the following has hold:
M
U.E
33
U.
E
1=
(14)
M
V
Σ
E
E
i=1
22

---

Published as a conference paper at ICLR 2022
for any 1 1 & E s.t. j' + j. Then, adding both inequalities yields, for all 1 i3s5 E
M
M
M
V
E jz
E
y/
i=
1
VI
M
|V]
0 E
U 34
E )
+
E
i=l
(15)
(E
EJA)
(E)) Ez,))
IM
V.
E
EJD
(E Ef.)
and the constraint of two vectors belonging to the same cluster j becomes
M
Σ 0.
(EE
EJ.
V
(16)
min IS
(o; (E) E...)
min
( (E, Ey)) T 2
1SJ'SE
Now, denote E(j) = -Ej. Our constraint satisfaction problem can be written as
min (u +') - EGi')) 2 0 18P'
(17)
By comparing Eq. 7 with Eq. 17, we observe that in our case, p is restricted to the set of vectors in R Therefore, we pick E TEE), where T(E) [w E V
E(j.i) EE(' i),0 s
Wi I 1, cos(U v' w) E [0, < < max || - loo S 1, which holds due to lp norm ordering and since is normalized in the Q scores expression. Adopting this notation simplifies the previous expression
The set EE) is non-empty if
min (v vd) MET(E)
(18)
Once again, recall that E is normalized. Therefore, we have
eij
Evy
+
Eij
evy
S d(i,t)
(19)
|eel12
||e,11/
||ell2
|Te,ll
which equivalently can be re-stated as (for e, 2 ej WLOG):
||ey| d(i
Cij
<
| lB
Gij ezj 1
Eij E.y.)
(20)
||e ||2
eg ezj <
(eg eys),
where take, example, d(i,i') = ||e;|la EEaj Egy:).
23

---

Published as a conference paper at ICLR 2022
The final expression for the sufficient condition for dynamics embeddings to belong to the same partition
min (0- v')'" T (ET(E)
(21)
s.1. (1() -((') S d(i,i')
for d(i,i') ||e/112 (Es- Eiy), which is similar to the Wasserstein-1 distance under d, i.e. (f(v, v').
We constructed an operator similar to F(d) in Ferns et al. (2004). can be computed by recursively applying F(d) at each v, v' E V pointwise, which is similar what done CTRL. This concludes our proof and shows how our clustering procedure can be viewed as finding reward-free bisimulations.
However, note that the exact interpretation of reward-free bisimulation relation depends on how is defined. Taking be two consecutive timesteps of state-action pairs yields the closest possible to the original definition of bisimulation, while sampling temporal keypoints far across the trajectory will induce different of properties.
24


Reviews:
[{"review_id": "XFXQhYKpopD", "reviewer": "Reviewer_z5qk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n1. This paper is well-written and easy to follow.\n2. The proposed method of training the encoder using clustering and prediction losses seems novel to me.\n3. Experimental results in all 16 tasks in the Procgen benchmark show that the proposed CTRL generally achieve zero-shot generalization.\n4. The additional analysis on clustering convergence and local perceptual changes are helpful for practical usage of CTRL.\n\nWeaknesses:\n1. A lot of details are omitted from the main paper, e.g. the detailed definitions of the clustering loss and the prediction loss.\n2. The proposed method does not significantly outperform baselines in many tasks. Other than bigfish, bossfight and starpilot, the improvement made by CTRL compared to baselines is relatively marginal.\n\nOther questions:\n1. Is it possible that the representation loss leads to a degenerate solution? For example, the encoder maps all observations to the same representation. I do not find a constraint to avoid degeneration.\n2. What is the relation between this proposed clustering+prediction loss and other contrastive learning based methods such as CURL? \n3. How does it compare to the Darla method[1]?\n\n[1] Higgins, Irina, et al. Darla: Improving zero-shot transfer in reinforcement learning.\n", "summary_of_the_paper": "This paper proposes a new zero-shot generalization method for RL that uses a novel clustering+prediction loss to learn a cross-trajectoriy encoder. The learned encoder is supposed to map behaviorally similar observations to similar representations. The algorithm generally outperforms representation RL baselines in experiment.", "summary_of_the_review": "This paper proposes a novel idea of learning representation with a clustering+prediction loss, which makes the representation generalize to novel observations from the same task distribution. The experimental results are convincing, although a little marginal in many environments. The proposed method has some potential to help representation learning in RL and zero-shot generalization.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "AU7DjNSs3Mj", "reviewer": "Reviewer_DS9R", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I found the paper interesting, but rather hard to understand what it really contributed and how several decisions came about. It is really only after I read MYOW and SwAV that I could understand several model architecture and loss choices, as the main text doesn’t cover them well enough and the Appendix is quite lacking. It also appeared to me that the proposed model is really close to MYOW and that fact wasn’t made as clear as I would have liked.\n\n\nAdditional comments and questions:\n\n1. The introduction and background are very clear and set up the problem well. The literature review is thorough and recent, and the overview of the algorithm in section 4 is also clear.\n2. My main issue was in understanding several model choices, when looking at Figure 1 especially. After looking at MYOW and SwAV, one can see how the overall model is basically MYOW with a different sequence of networks (a simpler one, which appears more plausible than MYOW to me, although see the next point) combined cleverly with SwAV’s soft-clustering model. CTRL also removes all multi-views aspects of these works, instead operating on trajectories directly.\n   1. It would have been much clearer to spell out these differences and what you did explicitly, given I would still consider that CTRL is significantly different from either of these two previous works.\n   2. One particular point which was confusing was the decision to make latent space be similar to adjacent clusters. This is taken directly from MYOW, but wasn’t clearly expressed here. It also appears slightly odd as there is no enforcement of a “global coverage” across the clusters (i.e. there is no extra repulsion term to far-away clusters), which to me would indicate that this method will have a chance to collapse (although one can be lucky by using target networks, ala BYOL).\n   3. Similarly, the choice to have a cascade of networks creating the two embedding for the clustering is rather surprising (and I would expect it to collapse to degenerate solutions given you’re not applying them in swapped ways like in SwAV).\n   4. The main text talks about trajectories, but T=2 if the Appendix is to be trusted. Hence in practice you are only clustering transitions. It would have been easier to present it that way (and the connection to bisimulation metrics would have been clearer).\n3. Many details aren’t in the main text, and the Appendix is equally light in details:\n   1. What is your encoder? Does it receive a history / have a recurrent state? Appendix 8.2 presents it as if it is instantaneous, which is quite a limitation.\n   2. How is FiLM used exactly? There are infinitely many ways to add it in.\n   3. It is very hard to follow the exact sequence of v/w/v’/w’ and how they are computed. One needs to carefully read SwAV to build an intuition, but this shouldn’t be that hard. No details are provided on what these are, apart from “\\psi_clust is a RNN” and “\\theta_clust is a MLP”\n   4. MYOW explicitly used predictors everywhere and were never regressing towards the latent used to cluster (i.e. `y = f_theta(x)` in their paper. See Figure 1, the augmented and clustering loss are applied on z and v, respectively.). Instead, you directly force `w’` to be predictable of `v’_c_i`. It would be valuable to discuss why you believe this is more appropriate (MYOW’s argue that this worked worst for them).\n   5. It would be quite valuable to bring back some of Section 8.2 into the main text, and add many more details to clarify what the model really is, I don’t believe this is reproducible in this current state.\n4. Table 1 is great and I would again like to flag that the number of baselines available are clearly to a great standard. This must have been hard to do well, so congratulations!\n   1. One caveat is that you are only reporting Evaluation performance on held out tasks (which is good, that was your metric), but I always would have appreciated seeing the performance on the training distribution. This would make it clearer if the effect is about improved generalization, or about the policies themselves being better or worse for particular baselines/ablations.\n5. Section 6.2 is quite anecdotal, and it isn’t clear that the effect discussed is that significant?\n6. Section 6.3 feels out of place and is quite underwhelming after having covered the results on Procgen. \n   1. It is a very peculiar toy situation, which isn’t that standard, and I’m not entirely sure what it’s trying to demonstrate in its current state. \n   2. Can you use the same silhouette score change to tune the number of clusters in Procgen too? Given it’s a post-hoc analysis, this doesn’t really appear like a real solution to this problem.\n7. Figure 4 in the Appendix has T=3 and T=5, but they show no difference? There is also no discussion of Figure 4, the section in 8.3 appears empty.\n8. As a small point, you keep referring to “views”, but this is overloaded and isn’t as appropriate as you do not actually have augmentations like in both previous work? Once again makes it hard to follow but should be easy to fix.\n9. You might want to add a citation to Van der Pol et al 2020 (https://arxiv.org/abs/2002.11963) for the bisimulation work, as they had a slightly different take than Zhang et al 2021.\n", "summary_of_the_paper": "This paper proposes an extension of the MYOW (Azabou et al 2021) self-supervised learning technique, combining it with SwAV (Caron et al 2021)’s method to perform online clustering, and adapt it for RL to assess generalisation performance on Procgen.\nIt compares against several recent baselines (DBC, PSE, CURL, Proto-RL and DIAYN), and outperforms them consistently on the tasks assessed (albeit by a small margin, depending on the reward scale of Procgen which I’m not extremely familiar with).", "summary_of_the_review": "Overall, I think this is nice paper, with good baselines and clear results, however, its presentation right now is quite lacking and hence I feel it may need some work before being ready for publication.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "4AmvbjfO9KC", "reviewer": "Reviewer_Kekc", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Originality and significance**:\nThis paper addresses zero-shot generalization, which is an important problem in RL. The idea of using unsupervised representation learning techniques to avoid overfitting in RL is not new. But the proposed method, CTRL, is somehow novel in that it uses both online clustering and cross-cluster prediction for representation learning, which has not been explored in the literature.\n\n**Quality**:\nThe overall quality of the paper is below the threshold for acceptance. Below are detailed comments.\n1. The setup for the ProcGen experiment presented in Section 6.1 is unclear. ProcGen has two difficulty modes and the main text is not clear about which one was used in the experiment. I guess the easy mode was used based on the number of training levels. Moreover, the 8 million steps training budget is unusual. The original ProcGen paper recommended 25 million steps for the easy mode. The choice of reporting test performance after 8 million steps of training seems arbitrary and needs further justification.\n2. I have some questions regarding the results in Table 1. Although the main text says \"...DAAC also exhibits good generalization performance...\", DAAC in fact performs worse than PPO (6 wins, 9 loses, and 1 tie in 16 games). These results conflict with the results reported in the original paper by Raileanu and Fergus. This is a bit surprising as Raileanu and Fergus opensourced their code. One possible explanation is that the discrepancy is due to the lower training budget, but it then leads back to the first question of why using 8 million steps rather than the recommended 25 million steps. Overall, the discrepancy between the reported results and the published results makes Table 1 less convincing.\n3. The experiment in Section 6.1 is very helpful. It verifies the importance of adapting the clusters online as the policy improves.\n4. I have some difficulty understanding the purpose of Section 6.3. The main text says \"To demonstrate the importance of identifying behavioral similarities...\" But after reading this section through, all I got was a way to select the hyperparameter for the number of clusters. I will be grateful if the authors can provide some clarification.\n\n**Clarity**:\nThe overall flow of the paper is easy to follow. But there are critical details missing in the method description which makes it hard to understand. Although the Sinkhorn-Knopp procedure and MYOW are existing methods, the authors should provide minimum details in the *main text* so that the readers can understand what is going on. For example, the concrete equations for computing $L_{\\text{clust}}$ and $L_{\\text{pred}}$ should be presented. Figure 1 can be explained in more detail. Currently, it contains unexplained notations like $\\textbf{e}_{c}$.\n\n**Minor issues & typos**:\n* The references seem truncated. For example, all the citations to Zhang _et al_ are not shown in the references.\n* The citation to Kendall _et al_ in the 3rd line of Section 1 should use \\citep.\n* The 7th line on page 3: \"bisimation\" --> \"bisimulation\".\n* The 4th last line on page 5: the citation to Raileanu and Fergus should use \\citep.\n* Bullet point 2 in Definition 1: if $c$ is a state, what does $s' \\in c$ mean? Is it a typo?\n* The line below Definition 1: the citation to Ferns _et al_ should use \\citet. ", "summary_of_the_paper": "This paper addresses the zero-shot generalization (ZSG) problem in reinforcement learning (RL). Specifically, it focuses on learning state representations that can generalize well. To this end, the authors propose a self-supervised learning method, cross trajectory representation learning (CTRL), that captures behaviorally similarity in the trajectory representations. The CTRL objective has two components, one for online clustering and the other one for cross-cluster prediction. CTRL is shown to be connected to bisimulation metrics in RL. The empirical study in the ProcGen benchmark shows that CTRL performs better than existing methods in the literature. Further experiments demonstrate the importance of adaptive online clustering and shed light on how to choose the number of clusters as a hyperparameter.", "summary_of_the_review": "I think this paper is mainly held back by two factors. First, critical details are missing in the main text which makes the proposed method hard to understand. Second, I have some questions regarding the main empirical results that may undermine the empirical significance of this work. Therefore, I don't think this paper is ready to publish yet.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "1b1d7hxeUY", "reviewer": "Reviewer_7TDG", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "#### **Strengths**\n- The paper tackles a very important problem in the deep RL community.\n- While various components that constitute CTRL were previously introduced in other works, I believe that the paper presents a novel combination of these methods and applies them in a novel problem setting.\n- The proposed method is simple and can be easily plugged in to solve ZSG problems other than Procgen.\n- The experiments are extensive, as the authors have compared CTRL to various baselines, including bisimulation metrics and SSL-based methods. The proposed method also outperforms quite significantly these baselines in most tasks.\n- The paper is well written and easy to understand.\n\n#### **Weaknesses**\n- The major issue with the paper is the lack of discussions and explanations on why the proposed objective functions allow the encoder to recognize the behavioral similarity. Consider a scenario where a robot arm needs to pick up a particular object on a table. Each time the robot performs its task, the surface of the table will randomly change its color, which means the distribution of tasks is now the distribution of the colors. In this imaginary scenario, maximizing $\\mathcal{L}_{\\text{clust}} + \\mathcal{L}_{\\text{pred}}$ might learn an encoder that only encodes the color information, which is not what we want to achieve. If we utilize the reward information, however, we can avoid this behavior by forcing the encoder to not encode what is not useful to predict the reward. So the reward information in this case is actually helpful for learning a good representation.\n- The visualization in Figure 8 in the Appendix shows that in practice CTRL does recognize the behavioral similarity. But again, I want to understand more on why this is the case, and why the encoder does not 'cheat' in maximizing the proposed objective function.\n\n#### **Minor comments**\n- In the first sentence of the second paragraph on Page 4, the distribution should be $P(\\mathcal{O}, \\mathcal{Z}, \\mathcal{S})$. Same problem for the following sentence.\n- 'repeatedly' was repeated twice in the caption for Figure 1.\n", "summary_of_the_paper": "The paper addresses the problem of zero short generalization (ZSG) in RL by improving the generalization of the encoder. A new self-supervised learning objective is proposed to encourage the encoder to map behaviorally similar observations to similar representations without the use of reward signals. The overall model is named Cross Trajectory Representation Learning (CTRL), which achieves better generalization performance on the challenging Procgen benchmark suite.", "summary_of_the_review": "In general, the paper is interesting. The proposed method is novel and while being simple, it works effectively and outperforms the baselines. However, I want the authors to discuss in more detail about why their method works well in practice while in theory it can fail.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Meta Review of Cross-Trajectory Representation Learning for Zero-Shot Generalization in RL

This work investigates a zero-shot generalization method for RL based on an online-clustering adapting it to RL. The intuition of this approach (called Cross Trajectory Representation Learning, CTRL) is that the self-supervised objective used will encourage the encoder to map behaviorally similar observations to similar representations without the use of reward signals. The authors performed experiments on the 16 procgen tasks, and compared it against several baselines (DBC, PSE, CURL, Proto-RL and DIAYN). The performance is generally better against baselines, but what I like about it is that a new approach to achieve such performance is proposed.

The scores were generally good (6, 6, 6, 3), and the 6's are overall positive with the work (both in the writing, breadth of experiments). Reviewer Kekc, who gave a score of 3, maintained their score, despite acknowledging the authors' responses. The main outstanding issue from Kekc is that they believe the paper should stick with the original 25M step protocol (with a larger training budget), rather than 8M steps. If that's the main issue for this paper to not be accepted into ICLR 2022, I feel this can be adequately addressed for the camera ready version. (Please note that while I disagree with the final score of 3 that Kekc gave, I find their review to be highly informative and useful, and would like to acknowledge Kekc for their input and discussion).

Based on the discussion and the reviews, and with the context behind the score of 3, I would like to be on the side of recommending this paper for acceptance, and urge strongly for the authors to conduct the 25M experiments as reviewer Kekc suggested (as Kekc also noted, the training curves are still going up, so just train them for a longer time). Even if the final results are not as good as the 8M, that's fine, just include them in the final camera ready version, since I believe this work to meet the bar, offers good insights into RL generalization, has a good breadth of experiments and baselines, and will be of great interest to the broader RL community.
Decision: Accept (Poster)