Original: C ROSS -T RAJECTORY R EPRESENTATION L EARNING FOR Z ERO -S HOT G ENERALIZATION IN RL
Revised: C ROSS -T RAJECTORY R EPRESENTATION L EARNING FOR Z ERO -S HOT G ENERALIZATION IN RL

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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 and Fergus (2021).
Revised: 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 and Fergus, 2021).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: Our experiments 1 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).

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Deep reinforcement learning (RL) has emerged as a powerful tool for building decision-making agents for domains with high-dimensional observation spaces, such as video games (Mnih et al., 2015), robotic manipulation (Levine et al., 2016), and autonomous driving (Kendall et al., 2019).
Revised: Deep reinforcement learning (RL) has emerged as a powerful tool for building decision-making agents for domains with high-dimensional observation spaces, such as video games (Mnih et al., 2015), robotic manipulation (Levine et al., 2016), and autonomous driving (Kendall et al., 2019).

Original: However, while deep RL agents may excel at 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 a robotic manipulator, driving under a variety of conditions, or coping with different levels in a game, remains a challenge.
Revised: However, while deep RL agents may excel at 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 a robotic manipulator, driving under a variety of conditions, or coping with different levels in a game, remains a challenge.

Original: 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.
Revised: 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.

Original: ZSG settings highlight the fact that generalization often cannot be solved by more training, as it can be too expensive or impossible to instantiate all possible real-world deployment scenarios a-priori.
Revised: ZSG settings highlight the fact that generalization often cannot be solved by more training, as it can be too expensive or impossible to instantiate all possible real-world deployment scenarios a-priori.

Original: In this work, we aim to improve ZSG in RL by proposing a new way of training the agent’s representation , a low-dimensional summary of information relevant to decision-making extracted from the agent’s high-dimensional observations.
Revised: In this work, we aim to improve ZSG in RL by proposing a new way of training the agent’s representation , a low-dimensional summary of information relevant to decision-making extracted from the agent’s high-dimensional observations.

Original: Outside of RL, representation learning can help 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 et al., 2020; Wu et al., 2020).
Revised: Outside of RL, representation learning can help 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 et al., 2020; Wu et al., 2020).

Original: In RL, unsupervised representation learning in the form of auxiliary objectives can be used to provide a richer learning signal over learning from reward alone, which helps the agent avoid overfitting on task-specific information (Raileanu and Fergus, 2021).
Revised: In RL, unsupervised representation learning in the form of auxiliary objectives can be used to provide a richer learning signal over learning from reward alone, which helps the agent avoid overfitting on task-specific information (Raileanu and Fergus, 2021).

Original: However, to our knowledge, no unsupervised learning method used in this way in RL has thus far been shown to substantially improve performance in ZSG over end-to-end reward-based methods (e.g., Cobbe et al., 2020).
Revised: However, to our knowledge, no unsupervised learning method used in this way in RL has thus far been shown to substantially improve performance in ZSG over end-to-end reward-based methods (e.g., Cobbe et al., 2020).

Original: We posit that using unsupervised (reward-free) learning to find representations that capture behavioral similarity across different trajectories will improve ZSG in RL.
Revised: We posit that using unsupervised (reward-free) learning to find representations that capture behavioral similarity across different trajectories will improve ZSG in RL.

Original: 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).
Revised: 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).

Original: We expand on this framework to improve ZSG performance, using unsupervised learning to train an agent that recognizes behavior similarity in a reward-free fashion.
Revised: We expand on this framework to improve ZSG performance, using unsupervised learning to train an agent that recognizes behavior similarity in a reward-free fashion.

Original: To do so, we propose Cross Trajectory Representation Learning (CTRL) , which applies a novel self-supervised learning (SSL) objective to pairs of trajectories drawn from the agent’s policies.
Revised: To do so, we propose Cross Trajectory Representation Learning (CTRL) , which applies a novel self-supervised learning (SSL) objective to pairs of trajectories drawn from the agent’s policies.

Original: For optimization, CTRL defines a prediction objective across trajectory representations from nearby partitions defined by an online clustering algorithm.
Revised: For optimization, CTRL defines a prediction objective across trajectory representations from nearby partitions defined by an online clustering algorithm.

Original: 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.
Revised: 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.

Original: Our main contributions are as follows:
Revised: Our main contributions are as follows:

Original: • 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 belief states.
Revised: • 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 belief states.

Original: 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.
Revised: 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.

Original: • We empirically show that CTRL improves zero-shot generalization in the challenging Procgen benchmark suite (Cobbe et al., 2020).
Revised: • We empirically show that CTRL improves zero-shot generalization in the challenging Procgen benchmark suite (Cobbe et al., 2020).

Original: Through a series of ablations, we highlight the importance of cross-trajectory views in boosting behavioral similarity.
Revised: Through a series of ablations, we highlight the importance of cross-trajectory views in boosting behavioral similarity.

Original: • We connect CTRL to the class of bisimulation methods, and provide sufficient conditions under which both formalisms can be equivalent.
Revised: • We connect CTRL to the class of bisimulation methods, and provide sufficient conditions under which both formalisms can be equivalent.

Original:  B ACKGROUND , MOTIVATION , AND RELATED WORKS
Revised:  B ACKGROUND , MOTIVATION , AND RELATED WORKS

Original: There are a broad class of ZSG settings in RL, such as generalization across reward functions (Barreto et al., 2016; Touati and Ollivier, 2021; Misra et al., 2020), observation spaces (Zhang et al., 2021;
Revised: There are a broad class of ZSG settings in RL, such as generalization across reward functions (Barreto et al., 2016; Touati and Ollivier, 2021; Misra et al., 2020), observation spaces (Zhang et al., 2021;

Original: Li et al., 2021; Raileanu and Fergus, 2021), or task dynamics (Rakelly et al., 2019).
Revised: Li et al., 2021; Raileanu and Fergus, 2021), or task dynamics (Rakelly et al., 2019).

Original: For each of these settings, there are a number of promising directions for improving ZSG performance: giving the agent better exploration policies (Van Roy and Wen, 2016;
Revised: For each of these settings, there are a number of promising directions for improving ZSG performance: giving the agent better exploration policies (Van Roy and Wen, 2016;

Original: 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).
Revised: 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).

Original: In this work, we focus on directly improving the agent’s representations .
Revised: In this work, we focus on directly improving the agent’s representations .

Original: The agent’s representations are high-level abstractions of observations or trajectories from the environment (e.g., the output of an encoder), and the desired property here is that one can easily learn a policy on top of that representation such that the combined model (i.e., the agent) generalizes to novel situations.
Revised: The agent’s representations are high-level abstractions of observations or trajectories from the environment (e.g., the output of an encoder), and the desired property here is that one can easily learn a policy on top of that representation such that the combined model (i.e., the agent) generalizes to novel situations.

Original: The tasks are assumed to share a 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.
Revised: The tasks are assumed to share a 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.

Original: 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 to readily do so.
Revised: 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 to readily do so.

Original: Unsupervised representation learning has been shown to improve generalization across domains, including zero-shot in vision (Sylvain et al., 2020; Wu et al., 2020) and sample-efficiency in RL (Ey- senbach et al., 2019; Schwarzer et al., 2021; Stooke et al., 2021).
Revised: Unsupervised representation learning has been shown to improve generalization across domains, including zero-shot in vision (Sylvain et al., 2020; Wu et al., 2020) and sample-efficiency in RL (Ey- senbach et al., 2019; Schwarzer et al., 2021; Stooke et al., 2021).

Original: 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.
Revised: 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.

Original: 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 ZSG performance.
Revised: 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 ZSG performance.

Original: Self-supervised learning and reinforcement learning.
Revised: Self-supervised learning and reinforcement learning.

Original: A successful class of models that incorporate unsupervised objectives to improve RL use self-supervised learning (SSL) (Anand et al., 2019;
Revised: A successful class of models that incorporate unsupervised objectives to improve RL use self-supervised learning (SSL) (Anand et al., 2019;

Original: Srinivas et al., 2020; Mazoure et al., 2020; Schwarzer et al., 2021; Stooke et al., 2021; Higgins et al., 2017).
Revised: Srinivas et al., 2020; Mazoure et al., 2020; Schwarzer et al., 2021; Stooke et al., 2021; Higgins et al., 2017).

Original: SSL formulates objectives by generating different views of the data, which are essentially transformed versions of the data, e.g., generated by using data augmentation or by sampling patches.
Revised: SSL formulates objectives by generating different views of the data, which are essentially transformed versions of the data, e.g., generated by using data augmentation or by sampling patches.

Original: While successful in their own way, prior works that combine SSL with RL do so by applying known
Revised: While successful in their own way, prior works that combine SSL with RL do so by applying known

Original: SSL algorithms (e.g., from vision, Hjelm et al., 2018; van den Oord et al., 2019; Bachman et al., 2019; Chen et al., 2020;
Revised: SSL algorithms (e.g., from vision, Hjelm et al., 2018; van den Oord et al., 2019; Bachman et al., 2019; Chen et al., 2020;

Original: He et al., 2020; Grill et al., 2020) to RL in a nearly off-the-shelf manner, predicting state representations within a given trajectory, only potentially using other trajectories as counterexamples in a contrastive loss.
Revised: He et al., 2020; Grill et al., 2020) to RL in a nearly off-the-shelf manner, predicting state representations within a given trajectory, only potentially using other trajectories as counterexamples in a contrastive loss.

Original: As such, these methods’ representations can have trouble generalizing latent behavioral patterns present in ostensibly different trajectories.
Revised: As such, these methods’ representations can have trouble generalizing latent behavioral patterns present in ostensibly different trajectories.

Original: Bisimulation metrics in reinforcement learning.
Revised: Bisimulation metrics in reinforcement learning.

Original: Our hypothesis is that ZSG is achievable if the agent recognizes behavioral similarity between trajectories based on their long-term evolution.
Revised: Our hypothesis is that ZSG is achievable if the agent recognizes behavioral similarity between trajectories based on their long-term evolution.

Original: Learning this sort of behavior similarity is a central characteristic of bisimulation metrics (Ferns et al., 2004), which assign a value of 0 to states which are behaviorally indistinguishable and have the same reward.
Revised: Learning this sort of behavior similarity is a central characteristic of bisimulation metrics (Ferns et al., 2004), which assign a value of 0 to states which are behaviorally indistinguishable and have the same reward.

Original: Reward-based bisimulation metrics have been shown to learn representations that have a number of useful properties, e.g.: smoothness (Gelada et al., 2019), visual invariance (Zhang et al., 2021), action equivariance (van der Pol et al., 2020) and multi-task adaptation (Zhang et al., 2020).
Revised: Reward-based bisimulation metrics have been shown to learn representations that have a number of useful properties, e.g.: smoothness (Gelada et al., 2019), visual invariance (Zhang et al., 2021), action equivariance (van der Pol et al., 2020) and multi-task adaptation (Zhang et al., 2020).

Original: For ZSG however, encoding relational information based on reward may not actually help (Misra et al., 2020;
Revised: For ZSG however, encoding relational information based on reward may not actually help (Misra et al., 2020;

Original: Touati and Ollivier, 2021; Yang and Nachum, 2021; Agarwal et al., 2021), as the agent may overfit to spurious correlations between high-dimensional observations and the reward signal seen during training.
Revised: Touati and Ollivier, 2021; Yang and Nachum, 2021; Agarwal et al., 2021), as the agent may overfit to spurious correlations between high-dimensional observations and the reward signal seen during training.

Original: HOMER (Misra et al., 2020) expands on the concept of bisimulation to learn behavioral similarity between states using unsupervised exploration at deployment.
Revised: HOMER (Misra et al., 2020) expands on the concept of bisimulation to learn behavioral similarity between states using unsupervised exploration at deployment.

Original: Among the existing methods,
Revised: Among the existing methods,

Original: PSEs (Agarwal et al., 2021) reward-free notion of behavioral similarity is conceptually the closest to
Revised: PSEs (Agarwal et al., 2021) reward-free notion of behavioral similarity is conceptually the closest to

Original: CTRL’s, and we compare these algorithms empirically in Section 6.
Revised: CTRL’s, and we compare these algorithms empirically in Section 6.

Original: However, algorithmically and in terms of their modes of operation, CTRL and PSE are very different.
Revised: However, algorithmically and in terms of their modes of operation, CTRL and PSE are very different.

Original: PSE assumes the availability of expert policies for training tasks and learns a representation using trajectories from these experts and an action distance measure, which it also assumes to be provided.
Revised: PSE assumes the availability of expert policies for training tasks and learns a representation using trajectories from these experts and an action distance measure, which it also assumes to be provided.

Original: CTRL doesn’t make these assumptions and learns a representation online from trajectories simultaneously generated by its substrate RL algorithm.
Revised: CTRL doesn’t make these assumptions and learns a representation online from trajectories simultaneously generated by its substrate RL algorithm.

Original: Mining views across unsupervised clusters.
Revised: Mining views across unsupervised clusters.

Original: Given our hypothesis that a model that learns behavioral 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.
Revised: Given our hypothesis that a model that learns behavioral 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.

Original: A simple and natural choice is to collect agent trajectories as examples of behaviors, then do clustering (online, similar to
Revised: A simple and natural choice is to collect agent trajectories as examples of behaviors, then do clustering (online, similar to

Original: Asano et al., 2020; Caron et al., 2020) over trajectories.
Revised: Asano et al., 2020; Caron et al., 2020) over trajectories.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: This would be counter-productive for our generalization goal.
Revised: This would be counter-productive for our generalization goal.

Original: We therefore use a second objective built on top of the structure provided by clustering to learn a more diverse set of similarities.
Revised: We therefore use a second objective built on top of the structure provided by clustering to learn a more diverse set of similarities.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original:  P ROBLEM STATEMENT AND PRELIMINARIES
Revised:  P ROBLEM STATEMENT AND PRELIMINARIES

Original: Formally, we define our problem setting w.r.t.
Revised: Formally, we define our problem setting w.r.t.

Original: a discrete-time Markov decision process (MDP)
Revised: a discrete-time Markov decision process (MDP)

Original: M ≜ ⟨S , A , P , R⟩ , where S is a state space, A is an action space, P : S ×
Revised: M ≜ ⟨S , A , P , R⟩ , where S is a state space, A is an action space, P : S ×

Original: A × S → [0 , 1] is a transition function characterizing environment dynamics, and R : S ×
Revised: A × S → [0 , 1] is a transition function characterizing environment dynamics, and R : S ×

Original: A → R is a reward function.
Revised: A → R is a reward function.

Original: M ’s 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.
Revised: M ’s 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.

Original: In practice, an agent usually receives observations but not the full information about the environment’s current state.
Revised: In practice, an agent usually receives observations but not the full information about the environment’s current state.

Original: Consider an observation space O and an observation function Z : S × O → [0 , 1] that define what observations an agent may receive and how these observations are generated (possibly stochastically) from M ’s states.
Revised: Consider an observation space O and an observation function Z : S × O → [0 , 1] that define what observations an agent may receive and how these observations are generated (possibly stochastically) from M ’s states.

Original: We define a task T as a partially observable MDP (POMDP) T = ⟨S , A , P , R , O , Z , s 0 ⟩ , where s 0 ∈ S is an initial state.
Revised: We define a task T as a partially observable MDP (POMDP) T = ⟨S , A , P , R , O , Z , s 0 ⟩ , where s 0 ∈ S is an initial state.

Original: Although many RL agents make decisions in a POMDP based only on the current observation o t or at most a few recent ones, in general this may require using information from the entire observation history o 1 , . . .
Revised: Although many RL agents make decisions in a POMDP based only on the current observation o t or at most a few recent ones, in general this may require using information from the entire observation history o 1 , . . .

Original: , o t so far.
Revised: , o t so far.

Original: Denoting the space of such histories as H , computing an agent’s behavior for task T amounts to finding a policy π : H ×
Revised: Denoting the space of such histories as H , computing an agent’s behavior for task T amounts to finding a policy π : H ×

Original: A → [0 , 1] with the optimal or near-optimal expected return from the initial state V πT ≜
Revised: A → [0 , 1] with the optimal or near-optimal expected return from the initial state V πT ≜

Original: E (cid:2)(cid:80) ∞ t =

Revised: E (cid:2)(cid:80) ∞ t =

Original: γ t R ( S t , π ( H t ) | S 

Revised: γ t R ( S t , π ( H t ) | S

Original: = s 0 (cid:3) , where
Revised: = s 0 (cid:3) , where

Original: S t and H t are random variables for the POMDP’s underlying state and agent’s observation history at time step t , respectively, and γ is a discount factor.
Revised: S t and H t are random variables for the POMDP’s underlying state and agent’s observation history at time step t , respectively, and γ is a discount factor.

Original: For an MDP M = ⟨S , A , P , R⟩ , a set O of observation spaces, and a set Z of observation functions w.r.t.
Revised: For an MDP M = ⟨S , A , P , R⟩ , a set O of observation spaces, and a set Z of observation functions w.r.t.

Original: S , let a task family be the POMDPs set
Revised: S , let a task family be the POMDPs set

Original: T M, O , Z ≜ {⟨S , A , P , R , O , Z , s 0 ⟩} O∈ O , Z∈ Z ,s 0 ∈S .
Revised: T M, O , Z ≜ {⟨S , A , P , R , O , Z , s 0 ⟩} O∈ O , Z∈ Z ,s 0 ∈S .

Original: We assume that different observation spaces in O have the same mathematical form, e.g., pixel tensors representing possible camera images, but correspond to qualitatively distinct subspaces of this larger space, such as subspaces of images depicting brightly and dimly lit scenes.
Revised: We assume that different observation spaces in O have the same mathematical form, e.g., pixel tensors representing possible camera images, but correspond to qualitatively distinct subspaces of this larger space, such as subspaces of images depicting brightly and dimly lit scenes.

Original: Our training and evaluation protocol is formalized w.r.t.
Revised: Our training and evaluation protocol is formalized w.r.t.

Original: a task distribution d ( T M, O , Z ) ≜ P ( O , Z , S ) over task family T M, O , Z , where P ( O , Z , S ) is a joint probability mass over observation spaces, observation functions, and initial states.
Revised: a task distribution d ( T M, O , Z ) ≜ P ( O , Z , S ) over task family T M, O , Z , where P ( O , Z , S ) is a joint probability mass over observation spaces, observation functions, and initial states.

Original: In the rest of the paper, M , O , and Z will be clear from context, and we will denote the task family as T and the task distribution as d ( T ) .
Revised: In the rest of the paper, M , O , and Z will be clear from context, and we will denote the task family as T and the task distribution as d ( T ) .

Original: For agent training, we choose N tasks from T , and denote the distribution d ( T ) restricted to these N tasks as d ( T N ) .
Revised: For agent training, we choose N tasks from T , and denote the distribution d ( T ) restricted to these N tasks as d ( T N ) .

Original: During the training phase, the RL agent learns via a series of epochs (themselves composed of episodes), by sampling a task from d ( T ) independently at the start of each episode, until the total number of time steps exceeds its training budget.
Revised: During the training phase, the RL agent learns via a series of epochs (themselves composed of episodes), by sampling a task from d ( T ) independently at the start of each episode, until the total number of time steps exceeds its training budget.

Original: 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 π .
Revised: 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 π .

Original: In representation learning-aided RL, a policy is viewed as a composition π = θ ◦ ϕ of an encoder ϕ :
Revised: In representation learning-aided RL, a policy is viewed as a composition π = θ ◦ ϕ of an encoder ϕ :

Original: H → E and a policy head θ : E × A → [0 , 1] , both of which are the outputs of the training phase.
Revised: H → E and a policy head θ : E × A → [0 , 1] , both of which are the outputs of the training phase.

Original: The training phase is followed by an evaluation phase, during which the policy is applied to tasks sampled from d ( T \ T N ) , distribution d restricted to the set of tasks T \ T N not seen during training.
Revised: The training phase is followed by an evaluation phase, during which the policy is applied to tasks sampled from d ( T \ T N ) , distribution d restricted to the set of tasks T \ T N not seen during training.

Original: Our focus on generalization means that we seek a policy π whose encoder ϕ allows it to maximize E T ∼ d ( T \ T N )
Revised: Our focus on generalization means that we seek a policy π whose encoder ϕ allows it to maximize E T ∼ d ( T \ T N )

Original: [ V πT ] despite being trained only on distribution d ( T N ) .
Revised: [ V πT ] despite being trained only on distribution d ( T N ) .

Original:  A LGORITHM
Revised:  A LGORITHM

Original: 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 RL-trained behavior prescribes a good action.
Revised: 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 RL-trained behavior prescribes a good action.

Original: 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.
Revised: 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.

Original: Like most online RL methods themselves, our algorithm operates in epochs, learning from a batch of trajectories in each epoch.
Revised: Like most online RL methods themselves, our algorithm operates in epochs, learning from a batch of trajectories in each epoch.

Original: CTRL assumes that all trajectories within each of its training batches come from the same policy.
Revised: CTRL assumes that all trajectories within each of its training batches come from the same policy.

Original: Before the RL algorithm updates the policy head in a given epoch, CTRL uses a trajectory batch from the current policy to update the encoder with gradients of a novel auxiliary loss L CTRL that we describe in this section.
Revised: Before the RL algorithm updates the policy head in a given epoch, CTRL uses a trajectory batch from the current policy to update the encoder with gradients of a novel auxiliary loss L CTRL that we describe in this section.

Original: 4. I NTUITION AND HIGH - LEVEL DESCRIPTION
Revised: 4. I NTUITION AND HIGH - LEVEL DESCRIPTION

Original: Algorithm overview.
Revised: Algorithm overview.

Original: For each trajectory batch, CTRL performs 4 operations:
Revised: For each trajectory batch, CTRL performs 4 operations:

Original: Apply the observation history (belief state) encoder ϕ to generate a low-dimensional rewardagnostic representation (“view”) of each trajectory.
Revised: Apply the observation history (belief state) encoder ϕ to generate a low-dimensional rewardagnostic representation (“view”) of each trajectory.

Original: Group trajectories’ views into C sets ( C is a tunable hyperparameter) using an online clustering algorithm with loss L clust (Equation (4), Section 4.2).
Revised: Group trajectories’ views into C sets ( C is a tunable hyperparameter) using an online clustering algorithm with loss L clust (Equation (4), Section 4.2).

Original: Using trajectory pairs selected from neighboring clusters, apply a predictive loss L pred (Equation (6), Section 4.2) to encourage ϕ to capture cross-cluster behavioral similarities.
Revised: Using trajectory pairs selected from neighboring clusters, apply a predictive loss L pred (Equation (6), Section 4.2) to encourage ϕ to capture cross-cluster behavioral similarities.

Original:  Update 
Revised: 4. Update ϕ with gradients of the total loss: L CTRL = L clust + L pred .

Original: ϕ with gradients of the total loss: L CTRL =

Original: L

Original: clust + L pred The schema in Figure 1 provides a high-level outline of these steps’ implementation and explains their interplay within CTRL, accompanied by an intuition for each step (below) a more detailed description in Section 4.2.
Revised: The schema in Figure 1 provides a high-level outline of these steps’ implementation and explains their interplay within CTRL, accompanied by an intuition for each step (below) a more detailed description in Section 4.2.

Original: We conduct ablations to show the effect of removing the clustering and predictive objectives of CTRL, with details in Appendix 8.3.
Revised: We conduct ablations to show the effect of removing the clustering and predictive objectives of CTRL, with details in Appendix 8.3.

Original: Clustering.
Revised: Clustering.

Original: C clusters can be viewed as corresponding to C latent “situations” in which an RL agent may find itself.
Revised: C clusters can be viewed as corresponding to C latent “situations” in which an RL agent may find itself.

Original: Each situation is essentially a group of belief states.
Revised: Each situation is essentially a group of belief states.

Original: 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.
Revised: 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.

Original: Under this hypothesis, an agent’s policy can be expected to produce C sets of roughly similar trajectories.
Revised: Under this hypothesis, an agent’s policy can be expected to produce C sets of roughly similar trajectories.

Original: CTRL’s clustering step (#2 above) is an attempt to recover these trajectory sets.
Revised: CTRL’s clustering step (#2 above) is an attempt to recover these trajectory sets.

Original: Since each trajectory consists of belief states, the purpose of
Revised: Since each trajectory consists of belief states, the purpose of

Original: L clust is to force the encoder ϕ to compute belief state representations that make trajectories within each cluster look similar in the latent space.
Revised: L clust is to force the encoder ϕ to compute belief state representations that make trajectories within each cluster look similar in the latent space.

Original: 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., Proto-RL, Yarats et al., 2021).
Revised: 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., Proto-RL, Yarats et al., 2021).

Original: Cross-cluster prediction.
Revised: Cross-cluster prediction.

Original: Note, however, that the clustering loss emphasizes the recognition of behavioral similarities within clusters.
Revised: Note, however, that the clustering loss emphasizes the recognition of behavioral similarities within clusters.

Original: This may hurt generalization, as the resulting centroids may not faithfully represent behaviors encountered at test time.
Revised: This may hurt generalization, as the resulting centroids may not faithfully represent behaviors encountered at test time.

Original: Our hypothesis is that encouraging encoder ϕ to induce latent-space similarities between trajectories from different but adjacent clusters will increase its ability to recognize behaviors in unseen test trajectories.
Revised: Our hypothesis is that encouraging encoder ϕ to induce latent-space similarities between trajectories from different but adjacent clusters will increase its ability to recognize behaviors in unseen test trajectories.

Original: While there are several ways to encourage cross-cluster representational similarity, using a mechanism similar to MYOW (Azabou et al., 2021) on trajectories drawn from neighboring clusters captures this idea particularly well.
Revised: While there are several ways to encourage cross-cluster representational similarity, using a mechanism similar to MYOW (Azabou et al., 2021) on trajectories drawn from neighboring clusters captures this idea particularly well.

Original: Namely, to get the cross-predictive loss L pred , we sample trajectory view pairs from neighboring clusters and apply the cosine-similarity loss to those pairs.
Revised: Namely, to get the cross-predictive loss L pred , we sample trajectory view pairs from neighboring clusters and apply the cosine-similarity loss to those pairs.

Original: Using reward guidance without reward signal for representation learning.
Revised: Using reward guidance without reward signal for representation learning.

Original: CTRL trains encoder ϕ only using the gradients of L CTRL ; the RL algorithm’s loss L RL trains only the policy head.
Revised: CTRL trains encoder ϕ only using the gradients of L CTRL ; the RL algorithm’s loss L RL trains only the policy head.

Original: Thus, encoder ϕ is isolated from the previously observed dangers of overfitting to the reward function (Raileanu and Fergus, 2021) that shapes L RL .
Revised: Thus, encoder ϕ is isolated from the previously observed dangers of overfitting to the reward function (Raileanu and Fergus, 2021) that shapes L RL .

Original: 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 L RL ’s gradients, which are reward-dependent.
Revised: 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 L RL ’s gradients, which are reward-dependent.

Original: 4. D ETAILS
Revised: 4. D ETAILS

Original: Below we describe the details of each of CTRL’s steps, with CTRL’s pseudocode presented in
Revised: Below we describe the details of each of CTRL’s steps, with CTRL’s pseudocode presented in

Original: Algorithm 1 in Appendix 8.1 .
Revised: Algorithm 1 in Appendix 8.1 .

Original: While in general the agent’s belief state at step t of a trajectory is the entire observation history h t =
Revised: While in general the agent’s belief state at step t of a trajectory is the entire observation history h t =

Original: ( o 1 , . . .
Revised: ( o 1 , . . .

Original: , o t ) , in the rest of the section we will assume h t =
Revised: , o t ) , in the rest of the section we will assume h t =

Original: ( o t ) and, in a slight abuse of notation, use ϕ ( o t ) instead of ϕ ( h t ) to simplify explanations 1 .
Revised: ( o t ) and, in a slight abuse of notation, use ϕ ( o t ) instead of ϕ ( h t ) to simplify explanations 2 .

Original: We emphasize, however, that CTRL equally applies in settings where the agent uses a much longer history of observations as its state.
Revised: We emphasize, however, that CTRL equally applies in settings where the agent uses a much longer history of observations as its state.

Original: In this case, ϕ would be recurrent or process stacks of frames.
Revised: In this case, ϕ would be recurrent or process stacks of frames.

Original: We also note that our CTRL implementation’s high-level algorithmic choices for clustering and cross-cluster prediction – Sinkhorn-Knopp and MYOW, respectively – come from prior works (Caron et al., 2020; Azabou et al., 2021).
Revised: We also note that our CTRL implementation’s high-level algorithmic choices for clustering and cross-cluster prediction – Sinkhorn-Knopp and MYOW, respectively – come from prior works (Caron et al., 2020; Azabou et al., 2021).

Original: Generating low-dimensional trajectory views with encoder ϕ .
Revised: Generating low-dimensional trajectory views with encoder ϕ .

Original: CTRL’s
Revised: CTRL’s

Original: input in each epoch is a trajectory batch { traj i } Bi =1 of size B .
Revised: input in each epoch is a trajectory batch { traj i } Bi =1 of size B .

Original: Assume all trajectories in the batch have the same length L .
Revised: Assume all trajectories in the batch have the same length L .

Original: For an integer hyperparameter T ≤ L , for each trajectory traj i =
Revised: For an integer hyperparameter T ≤ L , for each trajectory traj i =

Original: ( o 0 , a 0 , r 0 , . . .
Revised: ( o 0 , a 0 , r 0 , . . .

Original: o T b , a L , r L ) we independently and uniformly sample a subset of its steps τ i =
Revised: o T b , a L , r L ) we independently and uniformly sample a subset of its steps τ i =

Original: t 1 , . .
Revised: t 1 , . .

Original: . , t T to form a subtrajectory ( o t 1 , a t 1 , r t 1 , . . .
Revised: . , t T to form a subtrajectory ( o t 1 , a t 1 , r t 1 , . . .

Original: o t T , a t T , r t T ) .
Revised: o t T , a t T , r t T ) .

Original: We then encode this subtrajectory asu ( τ i ) i
Revised: We then encode this subtrajectory as

Original: [Equation]
Revised: [Equation]

Original: FiLM [Equation] where FiLM ( ϕ ( o t j ) , a t j ) is a common way of combining representations of different objects, akin to conditioning Perez et al.
Revised: where FiLM ( ϕ ( o t j ) , a t j ) is a common way of combining representations of different objects, akin to conditioning Perez et al.

Original: (2018), and the resulting vector u ( τ i )
Revised: (2018), and the resulting vector u ( τ i )

Original: i is in a low-dimensional space U .
Revised: i is in a low-dimensional space U .

Original: Note two aspects of the process of generating these vectors: (1) it drops rewards from the original trajectory and (2) it critically relies on ϕ whose parameter values are learned from previous epochs.
Revised: Note two aspects of the process of generating these vectors: (1) it drops rewards from the original trajectory and (2) it critically relies on ϕ whose parameter values are learned from previous epochs.

Original: Vectors u ( τ i ) i produced in this way are the trajectory views that the next steps of CTRL operate on.
Revised: Vectors u ( τ i ) i produced in this way are the trajectory views that the next steps of CTRL operate on.

Original: Clustering trajectory views.
Revised: Clustering trajectory views.

Original: CTRL groups trajectories from the epoch’s batch by clustering the set of their views { u ( τ i ) i } B i =1 .
Revised: CTRL groups trajectories from the epoch’s batch by clustering the set of their views { u ( τ i ) i } Bi =1 .

Original: 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., Proto-RL, Yarats et al., 2021).
Revised: 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., Proto-RL, Yarats et al., 2021).

Original: Since
Revised: Since

Original: CTRL operates online, Sinkhorn-Knopp is better-suited for the task than other clustering methods.
Revised: CTRL operates online, Sinkhorn-Knopp is better-suited for the task than other clustering methods.

Original: The clustering branch computes two views of each input u ( τ i ) i in a cascading fashion: first by passing it through a clustering encoder ψ clust : U → V , e.g. an RNN, to obtain a lower-dimensional view v i = ψ clust ( u ( τ i ) i ) , and then by passing v i through yet another network, an MLP θ clust : V → W , to produce view w i
Revised: The clustering branch computes two views of each input u ( τ i ) i in a cascading fashion: first by passing it through a clustering encoder ψ clust : U → V , e.g. an RNN, to obtain a lower-dimensional view v i = ψ clust ( u ( τ i ) i ) , and then by passing v i through yet another network, an MLP θ clust : V → W , to produce view w i

Original: =
Revised: =

Original: θ clust ( v i ) .
Revised: θ clust ( v i ) .

Original: The parameters of ψ clust and θ clust are learned through the epochs jointly with ϕ ’s. Like u ( τ i ) i , each v i and w i is a view of trajectory i ; the approach then consists in projecting v
Revised: The parameters of ψ clust and θ clust are learned through the epochs jointly with ϕ ’s. Like u ( τ i ) i , each v i and w i is a view of trajectory i ; the approach then consists in projecting v

Original: i ’s and w i ’s onto centroids of C clusters in two different ways and then computes a clustering loss that enforces consistency between v
Revised: i ’s and w i ’s onto centroids of C clusters in two different ways and then computes a clustering loss that enforces consistency between v

Original: i ’s and w i ’s cluster projections.
Revised: i ’s and w i ’s cluster projections.

Original: Specifically, we represent the centroid of each cluster c with a vector e c ∈ V , which are stacked into a matrix E .
Revised: Specifically, we represent the centroid of each cluster c with a vector e c ∈ V , which are stacked into a matrix E .

Original: These vectors are additional parameters in the joint optimization problem CTRL solves.
Revised: These vectors are additional parameters in the joint optimization problem CTRL solves.

Original: They can be regarded as views of C typical behaviors around which the trajectories’ views are regrouped.
Revised: They can be regarded as views of C typical behaviors around which the trajectories’ views are regrouped.

Original: To project trajectory i ’s v i views onto behavioral centroids learned from previous trajectory batches, CTRL computes a vector of soft assignments of v i to each centroid e c :
Revised: To project trajectory i ’s v i views onto behavioral centroids learned from previous trajectory batches, CTRL computes a vector of soft assignments of v i to each centroid e c :

Original: [Equation]
Revised: [Equation]

Original: and forms a B × C matrix Q whose i -th row is the soft assignment of v i .
Revised: and forms a B × C matrix Q whose i -th row is the soft assignment of v i .

Original: The resulting assignments may be very unbalanced, with most probability mass assigned to only a few clusters.
Revised: The resulting assignments may be very unbalanced, with most probability mass assigned to only a few clusters.

Original: Applying the
Revised: Applying the

Original: Sinkhorn-Knopp algorithm solves this issue by iteratively re-normalizing Q in order to obtain a more equal cluster membership (Cuturi, 2013), where the degree of re-normalization is controlled by a temperature parameter β .
Revised: Sinkhorn-Knopp algorithm solves this issue by iteratively re-normalizing Q in order to obtain a more equal cluster membership (Cuturi, 2013), where the degree of re-normalization is controlled by a temperature parameter β .

Original: The output of this operation is a matrix ˜ Q .
Revised: The output of this operation is a matrix ˜ Q .

Original: For each view w i (the projection of trajectory view v i ) CTRL computes the logarithm of its soft cluster assignments and treats these vectors as rows of another B × C matrix P :
Revised: For each view w i (the projection of trajectory view v i ) CTRL computes the logarithm of its soft cluster assignments and treats these vectors as rows of another B × C matrix P :

Original: [Equation]
Revised: [Equation]

Original: Finally, we compute the cross entropy between ˜ Q and P , which measures their inconsistency.
Revised: Finally, we compute the cross entropy between ˜ Q and P , which measures their inconsistency.

Original: This measure is taken as the clustering loss:
Revised: This measure is taken as the clustering loss:

Original: [Equation]
Revised: [Equation]

Original: Encouraging cross-cluster behavioral similarity.
Revised: Encouraging cross-cluster behavioral similarity.

Original: Note, however, that the clustering loss in the above step emphasizes the recognition of behavioral similarities within clusters.
Revised: Note, however, that the clustering loss in the above step emphasizes the recognition of behavioral similarities within clusters.

Original: This may hurt generalization, as the resulting centroids may not faithfully represent behaviors encountered at test time.
Revised: This may hurt generalization, as the resulting centroids may not faithfully represent behaviors encountered at test time.

Original: Our hypothesis is that encouraging encoder ϕ to induce latent-space similarities between trajectories from different but adjacent clusters will increase its ability to recognize behaviors in unseen test trajectories.
Revised: Our hypothesis is that encouraging encoder ϕ to induce latent-space similarities between trajectories from different but adjacent clusters will increase its ability to recognize behaviors in unseen test trajectories.

Original: While there are several ways to encourage cross-cluster representational similarity, using a mechanism similar to MYOW (Azabou et al., 2021) on trajectories drawn from neighboring clusters captures this idea particularly well.
Revised: While there are several ways to encourage cross-cluster representational similarity, using a mechanism similar to MYOW (Azabou et al., 2021) on trajectories drawn from neighboring clusters captures this idea particularly well.

Original: Namely, to get the cross-predictive loss L pred , we sample trajectory view pairs from neighboring clusters and apply the cosine-similarity loss to those pairs.
Revised: Namely, to get the cross-predictive loss L pred , we sample trajectory view pairs from neighboring clusters and apply the cosine-similarity loss to those pairs.

Original: To implement this idea, we define a measure of cluster proximity via a matrix D of cosine similarities between cluster centroids: for clusters k and l , D kl =
Revised: To implement this idea, we define a measure of cluster proximity via a matrix D of cosine similarities between cluster centroids: for clusters k and l , D kl =

Original: ||
Revised: ||

Original: e k − e l || 

Revised: e k − e l ||

Original: (Grill et al., 2020).
Revised: (Grill et al., 2020).

Original: Recall that in the previous step, the clustering branch computed a matrix Q whose rows Q
Revised: Recall that in the previous step, the clustering branch computed a matrix Q whose rows Q

Original: i are soft assignments of trajectory
Revised: i are soft assignments of trajectory

Original: i ’s view v i to clusters.
Revised: i ’s view v i to clusters.

Original: In this step, we convert these soft assignments to hard ones by associating a trajectory’s view v i with cluster c i = argmax 1 ≤ c ′ ≤ C ˜ Q i and treating a cluster c as consisting of trajectories with indices in the set T c
Revised: In this step, we convert these soft assignments to hard ones by associating a trajectory’s view v i with cluster c i = argmax 1 ≤ c ′ ≤ C ˜ Q i and treating a cluster c as consisting of trajectories with indices in the set T c

Original: =
Revised: =

Original: { i | c = argmax 1 ≤ c ′ ≤ C ˜ Q i } .
Revised: { i | c = argmax 1 ≤ c ′ ≤ C ˜ Q i } .

Original: To assess how predictive a trajectory embedding u i is of a trajectory embedding u j , like in the clustering step we will use two special helper maps, ψ pred : U → V to obtain a reduced-dimensionality view v ′ =
Revised: To assess how predictive a trajectory embedding u i is of a trajectory embedding u j , like in the clustering step we will use two special helper maps, ψ pred : U → V to obtain a reduced-dimensionality view v ′ =

Original: ψ pred ( u ) and θ pred to further project v ′ to w ′ =
Revised: ψ pred ( u ) and θ pred to further project v ′ to w ′ =

Original: θ pred ( v ′ ) .
Revised: θ pred ( v ′ ) .

Original: CTRL proceeds by repeatedly sampling trajectories, which we call anchor trajectories , from the batch, with their associated embeddings u .
Revised: CTRL proceeds by repeatedly sampling trajectories, which we call anchor trajectories , from the batch, with their associated embeddings u .

Original: For each anchor trajectory n , consider K clusters c 1 , . .
Revised: For each anchor trajectory n , consider K clusters c 1 , . .

Original: c K nearest to n ’s cluster c n , as defined by the indices of K largest values in row c n of matrix D (we exclude c n itself when determining c n ’s nearest clusters).
Revised: c K nearest to n ’s cluster c n , as defined by the indices of K largest values in row c n of matrix D (we exclude c n itself when determining c n ’s nearest clusters).

Original: Borrowing ideas from the MYOW approach
Revised: Borrowing ideas from the MYOW approach

Original: Azabou et al.
Revised: Azabou et al.

Original: (2021), CTRL mines a view for u n by randomly choosing a trajectory with embedding u ( n ) c k from each of the neighboring clusters and computing its view v ′ c k = ψ pred ( u ( n ) c k ) .
Revised: (2021), CTRL mines a view for u n by randomly choosing a trajectory with embedding u ( n ) c k from each of the neighboring clusters and computing its view v ′ c k = ψ pred ( u ( n ) c k ) .

Original: We call the neighbors’ views v ′ c 1 , . .
Revised: We call the neighbors’ views v ′ c 1 , . .

Original: . , v ′ c K trajectory n ’s mined views .
Revised: . , v ′ c K trajectory n ’s mined views .

Original: For the final operation in this step, CTRL computes trajectory n ’s predictive view w ′ n = θ pred ( ψ pred ( u n )) and measures the distance from it to trajectory n ’s mined views: 
Revised: For the final operation in this step, CTRL computes trajectory n ’s predictive view w ′ n = θ pred ( ψ pred ( u n )) and measures the distance from it to trajectory n ’s mined views: [Equation]

Original: [Equation]
Revised: [Equation]

Original: N regulates the number of anchor trajectories to be sampled, so the total prediction loss is
Revised: N regulates the number of anchor trajectories to be sampled, so the total prediction loss is

Original: [Equation]
Revised: [Equation]

Original: Updating encoder ϕ using reward guidance without reward signal for representation learning.
Revised: Updating encoder ϕ using reward guidance without reward signal for representation learning.

Original: Note that CTRL’s total loss L CTRL depends on the parameters of encoder ϕ as well as of clustering networks ϕ clust and θ clust , prediction networks ϕ clust and θ clust , and cluster centroids e c , 1 ≤ c ≤ C .
Revised: Note that CTRL’s total loss L CTRL depends on the parameters of encoder ϕ as well as of clustering networks ϕ clust and θ clust , prediction networks ϕ clust and θ clust , and cluster centroids e c , 1 ≤ c ≤ C .

Original: In each epoch, CTRL updates all these parameters to minimize L CTRL .
Revised: In each epoch, CTRL updates all these parameters to minimize L CTRL .

Original:  C ONNECTION TO BISIMULATION
Revised:  C ONNECTION TO BISIMULATION

Original: Deep bisimulation metrics are tightly connected to the underlying mechanism of mining behaviorally similar trajectories of CTRL.
Revised: Deep bisimulation metrics are tightly connected to the underlying mechanism of mining behaviorally similar trajectories of CTRL.

Original: They operate on a latent-dimensional space and, as is the case for
Revised: They operate on a latent-dimensional space and, as is the case for

Original: 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.
Revised: 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.

Original: In this section, we aim to highlight a functional similarity between bisimulation metrics and CTRL.
Revised: In this section, we aim to highlight a functional similarity between bisimulation metrics and CTRL.

Original: Definition 1 A bisimilation relation E ⊆ S × S is a binary relation which satisfies, ∀ ( s, t ) ∈
Revised: Definition 1 A bisimilation relation E ⊆ S × S is a binary relation which satisfies, ∀ ( s, t ) ∈

Original: E :
Revised: E :

Original: [Equation]
Revised: [Equation]

Original: In practice, rewards and transition probabilities rarely match exactly.
Revised: In practice, rewards and transition probabilities rarely match exactly.

Original: For this reason, Ferns et al.
Revised: For this reason, Ferns et al.

Original: (2004) proposed a smooth alternative to bisimulation relations in the form of bisimulation metrics, which can be found by solving a recursive equation involving the Wasserstein-1 distance W 1 between transition probabilities.
Revised: (2004) proposed a smooth alternative to bisimulation relations in the form of bisimulation metrics, which can be found by solving a recursive equation involving the Wasserstein-1 distance W 1 between transition probabilities.

Original: W 1 can be found by solving the following linear programming (Villani, 2008), where we let Γ = { v ∈ R |V| : 0 ≤ v i ≤ 1 ∀ 1 ≤ i ≤ |V|
Revised: W 1 can be found by solving the following linear programming (Villani, 2008), where we let Γ = { v ∈ R |V| : 0 ≤ v i ≤ 1 ∀ 1 ≤ i ≤ |V|} :

Original: [Equation]
Revised: [Equation]

Original: where µ is a vector whose elements are constrained between 0 and 1.
Revised: where µ is a vector whose elements are constrained between 0 and 1.

Original: In practice, bisimulation metrics are used to enforce a temporal continuity of the latent space by minimization of the W 1 loss between training state-action pairs.
Revised: In practice, bisimulation metrics are used to enforce a temporal continuity of the latent space by minimization of the W 1 loss between training state-action pairs.

Original: Therefore, to show a connection of CTRL to (reward-free) bisimulation metrics, it is sufficient to show that two trajectories are mapped to the same partition if their induced
Revised: Therefore, to show a connection of CTRL to (reward-free) bisimulation metrics, it is sufficient to show that two trajectories are mapped to the same partition if their induced

Original: W 1 distance is arbitrarily small.
Revised: W 1 distance is arbitrarily small.

Original: In our (informal) argument that follows, we assume that CTRL samples two consecutive timesteps and encodes them into v ; the exact form of v dictates the nature of the behavioral similarity.
Revised: In our (informal) argument that follows, we assume that CTRL samples two consecutive timesteps and encodes them into v ; the exact form of v dictates the nature of the behavioral similarity.

Original: The proof can be found in Appendix 8.7.
Revised: The proof can be found in Appendix 8.7.

Original: Proposition 1 (Informal) Let M be an MDP where R ( s, a )
Revised: Proposition 1 (Informal) Let M be an MDP where R ( s, a )

Original: = 

Revised: =

Original: for all ( s, a ) ∈ S × A and let v , v ′ ∈ V be two dynamics embeddings in M .
Revised: for all ( s, a ) ∈ S × A and let v , v ′ ∈ V be two dynamics embeddings in M .

Original: The clustering operation between v , v ′ induces a reward-free bisimilarity metric W 1 ( P [ v ] , P [ v ′ ]) between induced distributions P [ v ] and P [ v ′ ] .
Revised: The clustering operation between v , v ′ induces a reward-free bisimilarity metric W 1 ( P [ v ] , P [ v ′ ]) between induced distributions P [ v ] and P [ v ′ ] .

Original:  E MPIRICAL EVALUATION
Revised:  E MPIRICAL EVALUATION

Original: We compare CTRL against strong RL baselines: DAAC (Raileanu and Fergus, 2021) – the current state-of-the-art on the challenging generalization benchmark suite Procgen (Cobbe et al., 2020), and PPO (Schulman et al., 2017).
Revised: We compare CTRL against strong RL baselines: DAAC (Raileanu and Fergus, 2021) – the current state-of-the-art on the challenging generalization benchmark suite Procgen (Cobbe et al., 2020), and PPO (Schulman et al., 2017).

Original: 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.
Revised: 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.

Original: We then compare to several unsupervised and SSL auxiliary objectives used in conjunction with PPO.
Revised: We then compare to several unsupervised and SSL auxiliary objectives used in conjunction with PPO.

Original: DIAYN (Eysenbach et al., 2019) is an unsupervised skillbased exploration method which we adapt to the online setting by uniformly sampling skills.
Revised: DIAYN (Eysenbach et al., 2019) is an unsupervised skillbased exploration method which we adapt to the online setting by uniformly sampling skills.

Original: Its notion of skills has some similarities to the notion of clusters in CTRL.
Revised: Its notion of skills has some similarities to the notion of clusters in CTRL.

Original: We also compare with two SSL-based auxiliary objectives: CURL (Srinivas et al., 2020), a common SSL baseline which contrasts augmented instances of the same state, and Proto-RL (Yarats et al., 2021), which we adapt for this generalization setting.
Revised: We also compare with two SSL-based auxiliary objectives: CURL (Srinivas et al., 2020), a common SSL baseline which contrasts augmented instances of the same state, and Proto-RL (Yarats et al., 2021), which we adapt for this generalization setting.

Original: 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).
Revised: 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).

Original: PSE assumes policies for training tasks to be given and, like Agarwal et al.
Revised: PSE assumes policies for training tasks to be given and, like Agarwal et al.

Original: (2021), we ran it both with random and high-quality policies pretrained with extra computation budget.
Revised: (2021), we ran it both with random and high-quality policies pretrained with extra computation budget.

Original: See Appendix 8.2 for details.
Revised: See Appendix 8.2 for details.

Original: The Procgen benchmark suite, which we use in our experiments, consists of 16 video games (see
Revised: The Procgen benchmark suite, which we use in our experiments, consists of 16 video games (see

Original: Table 1).
Revised: Table 1).

Original: Procgen procedurally generates distinct levels for each game.
Revised: Procgen procedurally generates distinct levels for each game.

Original: The number of levels for each game is virtually unlimited.
Revised: The number of levels for each game is virtually unlimited.

Original: Levels within a game which share common game rules and objectives but differ in level design such as the number of projectiles, background colors, item placements throughout the level and other game assets.
Revised: Levels within a game which share common game rules and objectives but differ in level design such as the number of projectiles, background colors, item placements throughout the level and other game assets.

Original: All of this makes Procgen a suitable benchmark for zero-shot generalization.
Revised: All of this makes Procgen a suitable benchmark for zero-shot generalization.

Original: Using our notation from Section 3, for each of 16 games, we train on a uniform distribution d ( T N ) over N = 200 “easy” levels of the game and evaluate on d ( T \ T N ) , i.e., a uniform distribution over the game’s “easy” levels not seen during training.
Revised: Using our notation from Section 3, for each of 16 games, we train on a uniform distribution d ( T N ) over N = 200 “easy” levels of the game and evaluate on d ( T \ T N ) , i.e., a uniform distribution over the game’s “easy” levels not seen during training.

Original: Following Mohanty et al.
Revised: Following Mohanty et al.

Original: (2021), we report results after 8M steps of training, since this demonstrates the quality of ZSG that various representation learning methods can achieve quickly .
Revised: (2021), we report results after 8M steps of training, since this demonstrates the quality of ZSG that various representation learning methods can achieve quickly .

Original: However, Figure 4 in Appendix 8.
Revised: However, Figure 4 in Appendix 8.

Original: also provides results after 25M steps of training, as in the original Procgen paper (Cobbe et al., 2020).
Revised: also provides results after 25M steps of training, as in the original Procgen paper (Cobbe et al., 2020).

Original: [Table]
Revised: [Table]

Original: Main results.
Revised: Main results.

Original: As Table 1 shows, PPO+CTRL outperforms all other baselines, including DAAC, on most games.
Revised: As Table 1 shows, PPO+CTRL outperforms all other baselines, including DAAC, on most games.

Original: Notably, bisimulation-based approaches other than CTRL– DBC as well as PSE with both random and expert data-gathering policies – exhibit lower gains than others.
Revised: Notably, bisimulation-based approaches other than CTRL– DBC as well as PSE with both random and expert data-gathering policies – exhibit lower gains than others.

Original: While this can be surprising, recent work has seen similar results when applying DBC to tasks with unseen backgrounds (Li et al., 2021).
Revised: While this can be surprising, recent work has seen similar results when applying DBC to tasks with unseen backgrounds (Li et al., 2021).

Original: PSE’s inferior performance may be due to the policy similarity metric, which PSE requires as input and which we took from Agarwal et al.
Revised: PSE’s inferior performance may be due to the policy similarity metric, which PSE requires as input and which we took from Agarwal et al.

Original: (2021), being poorly suited to
Revised: (2021), being poorly suited to

Original: Procgen.
Revised: Procgen.

Original: This highlights an important difference between CTRL and PSE: CTRL doesn’t need a policy similarity metric, since it implicitly induces such a metric based on trajectory “signatures”.
Revised: This highlights an important difference between CTRL and PSE: CTRL doesn’t need a policy similarity metric, since it implicitly induces such a metric based on trajectory “signatures”.

Original: Despite training static prototypes for 8M timesteps and adapting the RL head for 8M additional timesteps (see Appendix 8.2 for details), Proto-RL performs worse than CTRL.
Revised: Despite training static prototypes for 8M timesteps and adapting the RL head for 8M additional timesteps (see Appendix 8.2 for details), Proto-RL performs worse than CTRL.

Original: This suggests that the temporal aspect of clustering is key for ZSG, a hypothesis we explore further in Section 8.4.
Revised: This suggests that the temporal aspect of clustering is key for ZSG, a hypothesis we explore further in Section 8.4.

Original: 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.
Revised: 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.

Original: DAAC also exhibits good generalization performance, but inherits from PPG (Cobbe et al., 2021)
Revised: DAAC also exhibits good generalization performance, but inherits from PPG (Cobbe et al., 2021)

Original: the separation of the value and policy branches both parameter-wise and by introducing distinct training phases, an overhead which CTRL manages to avoid.
Revised: the separation of the value and policy branches both parameter-wise and by introducing distinct training phases, an overhead which CTRL manages to avoid.

Original: 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 a toy example how learning behavioral similarities captures local perceptual changes (Appendix 8.5).
Revised: 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 a toy example how learning behavioral similarities captures local perceptual changes (Appendix 8.5).

Original:  D ISCUSSION
Revised:  C ONCLUSIONS

Original: In this work, we proposed CTRL, a novel representation learning algorithmfor RL that encourages zero-shot generalization across pixel-based observation spaces.
Revised: This work proposed CTRL, a novel representation learning algorithm that facilitates zero-shot generalization of RL policies in high-dimensional observation spaces.

Original: CTRLoperates by grouping behaviorally similar trajectories together, which can be viewed as inducing an unsupervised reward-agnostic bisimulation metric over observation histories, learned over transitions encountered by policies from an RL algorithm’s value improvement path (Dabney et al., 2021). 
Revised: CTRL can be viewed as inducing an unsupervised reward-agnostic bisimulation metric over observation histories, learned over transitions encountered by policies from an RL algorithm’s value improvement path (Dabney et al., 2021). We hope that in the future CTRL will inspire other representation learning methods based on capturing belief states’s behavioral similarity, which will be capable of policy generalization across greater variations in environment dynamics.

Original: update
Revised: update

Original: RL
Revised: RL

Original: parameters 8. E XPERIMENT DETAILS
Revised: parameters 8. E XPERIMENT DETAILS

Original: We implemented all algorithms on top of the IMPALA architecture, which was shown to perform well on Procgen (Cobbe et al., 2020).
Revised: We implemented all algorithms on top of the IMPALA architecture, which was shown to perform well on Procgen (Cobbe et al., 2020).

Original: In the Procgen experiments, Proto-RL was ran without any intrinsic rewards (since the domains are not exploration-focused) by first jointly training the representation and RL losses for 8M timesteps, after which only the RL loss was optimized for an additional 8M steps (16M steps in total).
Revised: In the Procgen experiments, Proto-RL was ran without any intrinsic rewards (since the domains are not exploration-focused) by first jointly training the representation and RL losses for 8M timesteps, after which only the RL loss was optimized for an additional 8M steps (16M steps in total).

Original: Similarly, DIAYN was also run with the pre-training phase of 8M and then
Revised: Similarly, DIAYN was also run with the pre-training phase of 8M and then

Original: RL only objective for the second 8M phase.
Revised: RL only objective for the second 8M phase.

Original: Like Proto-RL and DIAYN, PSE needed extra training budget and extra adjustments for a fair comparison to CTRL (see Section 2).
Revised: Like Proto-RL and DIAYN, PSE needed extra training budget and extra adjustments for a fair comparison to CTRL (see Section 2).

Original: Agarwal et al.
Revised: Agarwal et al.

Original: (2021) used PSE only with Soft Actor-Critic, which doesn’t perform well on Procgen.
Revised: (2021) used PSE only with Soft Actor-Critic, which doesn’t perform well on Procgen.

Original: Therefore, we carried over PSE’ available implementation
Revised: Therefore, we carried over PSE’ available implementation

Original: [Table]
Revised: [Table]

Original: from https://agarwl.github.io/pse/ to our codebase, with the help of PSE’ authors, to combine it with the same PPO implementation that CTRL used.
Revised: from https://agarwl.github.io/pse/ to our codebase, with the help of PSE’ authors, to combine it with the same PPO implementation that CTRL used.

Original: PSE assumes being given policies for training problem instances (Procgen levels).
Revised: PSE assumes being given policies for training problem instances (Procgen levels).

Original: Like Agarwal et al.
Revised: Like Agarwal et al.

Original: (2021), we ran PSE using both random and high-quality pretrained policies for these problems.
Revised: (2021), we ran PSE using both random and high-quality pretrained policies for these problems.

Original: In the latter case, we pretrained expert policies for the first 40 levels of Procgen to generate training trajectories for PSE.
Revised: In the latter case, we pretrained expert policies for the first 40 levels of Procgen to generate training trajectories for PSE.

Original: Each level’s expert was trained on 0.5M environment steps.
Revised: Each level’s expert was trained on 0.5M environment steps.

Original: Note that this is much less than 8M steps we used for policy training in experiments with other algorithms, but this is because each expert needed to be good only for a single level , and we verified that they indeed were.
Revised: Note that this is much less than 8M steps we used for policy training in experiments with other algorithms, but this is because each expert needed to be good only for a single level , and we verified that they indeed were.

Original: 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.
Revised: 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.

Original: We don’t believe more than 40 experts per game would have made a difference.
Revised: We don’t believe more than 40 experts per game would have made a difference.

Original: Given policies for the training levels, PPO+PSE’s training on Procgen mimicked CTRL’s:
Revised: Given policies for the training levels, PPO+PSE’s training on Procgen mimicked CTRL’s:

Original: PPO+PSE trained by interacting with the first 200 levels for 8M steps and was evaluated on the rest.
Revised: PPO+PSE trained by interacting with the first 200 levels for 8M steps and was evaluated on the rest.

Original: However, during the training, PPO+PSE sampled an additional 1M interactions from the pretrained traiing-level policies.
Revised: However, during the training, PPO+PSE sampled an additional 1M interactions from the pretrained traiing-level policies.

Original: the process was repeated for all 16 games, for 10 seeds each.
Revised: the process was repeated for all 16 games, for 10 seeds each.

Original: We did hyperparameter grid search on PPO+PSE’s hyperparameters for PSE – loss coefficient values of (0 . 1 , 1 , 2) and temperature (0 . 1 , 0 .
Revised: We did hyperparameter grid search on PPO+PSE’s hyperparameters for PSE – loss coefficient values of (0 . 1 , 1 , 2) and temperature (0 . 1 , 0 .

Original: PPO’s hyperparameters were the same as in CTRL.
Revised: PPO’s hyperparameters were the same as in CTRL.

Original: Thus, due to the need to pretrain and gather data with per-level expert policies, PPO+PSE received 0 .
Revised: Thus, due to the need to pretrain and gather data with per-level expert policies, PPO+PSE received 0 .

Original: 5 · 20 M + 1 M = 

Revised: 5 · 20 M + 1 M =

Original: M extra environment interactions compared to CTRL, i.e., used 21 M/ 8 M = 2 .
Revised: M extra environment interactions compared to CTRL, i.e., used 21 M/ 8 M = 2 .

Original: 6 × more training data than the latter.
Revised: 6 × more training data than the latter.

Original: 8. A DDITIONAL RESULTSZSG over 25M training steps and policy performance throughout training
Revised: 8. A DDITIONAL RESULTS

Revised: 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 2 shows the training performance of agents on 8M frames, Figure 3 the test performance of agents on 8M frames and all levels, and, finally, Figure 4 shows the test performance of agents on 25M frames and all levels.

Original: Ablations on algorithm components We ran multiple versions of the algorithm to identify the key components which make CTRL perform well in Procgen.
Revised: Ablations on algorithm components We ran multiple versions of the algorithm to identify the key components which make CTRL perform well in Procgen.

Original: The first modification, CTRL consecutive
Revised: The first modification, CTRL consecutive

Original: T consists in running our algorithm but sampling consecutive timesteps, that is t
Revised: T consists in running our algorithm but sampling consecutive timesteps, that is t

Original: i +

Revised: i +

Original: =
Revised: =

Original: t 1 + i for all tand 0 ≤ i ≤ T .
Revised: t 1 + i for all tand 0 ≤ i ≤ T .

Original: The second modification, CTRL no action removes the action conditioning layer in the log-softmax probability p t and in the cluster scores q t , to test the importance of action information for cluster membership prediction.
Revised: The second modification, CTRL no action removes the action conditioning layer in the log-softmax probability p t and in the cluster scores q t , to test the importance of action information for cluster membership prediction.

Original: The third modification, CTRL no cluster removes the clustering loss, and only restricts to mining and predicting nearby neighbors in the batch.
Revised: The third modification, CTRL no cluster removes the clustering loss, and only restricts to mining and predicting nearby neighbors in the batch.

Original: Finally, the last modification, CTRL no pred removes the loss predicting samples from neighboring partitions and only relies on the clustering loss to update its representation.
Revised: Finally, the last modification, CTRL no pred removes the loss predicting samples from neighboring partitions and only relies on the clustering loss to update its representation.

Original: 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[Table] up on the local dynamics present in the MDP and (3) both clustering and predictive objectives are essential to the good performance of our algorithm.
Revised: 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 to the good performance of our algorithm.

Original: Results of the last column are computed overseeds, rest over 3 seeds.
Revised: Results of the last column are computed overseeds, rest over 3 seeds.

Original: Ablation on number of clusters and clustering timesteps How should one determine the optimal number of clusters in a complex domain?
Revised: Ablation on number of clusters and clustering timesteps How should one determine the optimal number of clusters in a complex domain?

Original: Can the number of clusters be chosen a priori running any training?
Revised: Can the number of clusters be chosen a priori running any training?

Original: Below, we provide some partial answers to these questions.
Revised: Below, we provide some partial answers to these questions.

Original: First, the optimal (or true) number of clusters is domain-specific, as it depends on the exact connectivity structure of the MDP at hand.
Revised: First, the optimal (or true) number of clusters is domain-specific, as it depends on the exact connectivity structure of the MDP at hand.

Original: Second, the length of the clusters, i.e. the number of trajectory timesteps passed to Sinkhorn-Knopp can widely impact the nature of learned representations, and hence the downstream performance of the agent.
Revised: Second, the length of the clusters, i.e. the number of trajectory timesteps passed to Sinkhorn-Knopp can widely impact the nature of learned representations, and hence the downstream performance of the agent.

Original: Temporal connectivity of the clusters Are clusters consistent in time for a given trajectory?
Revised: Temporal connectivity of the clusters Are clusters consistent in time for a given trajectory?

Original: To verify this, we trained CTRL on 1 million frames of the bigfish game.
Revised: To verify this, we trained CTRL on 1 million frames of the bigfish game.

Original: For every T states in a given trajectory, we have computed the hard cluster assignment to the nearest cluster, which yields  a sequence of partitions.
Revised: For every T states in a given trajectory, we have computed the hard cluster assignment to the nearest cluster, which yields [Table] a sequence of partitions.

Original: We then computed the cosine similarity between time-adjacent cluster centroids, the metric reported on the smoothed graph below.
Revised: We then computed the cosine similarity between time-adjacent cluster centroids, the metric reported on the smoothed graph below [Table]

Original: Loss landscape of L clust and L pred
Revised: Loss landscape of L clust and L pred

Original: Works relying on 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.
Revised: Works relying on 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.

Original: For example, (Gelada et al., 2019) showed how their dynamics and reward losses are inversely proportional to each other early on in the training, taking a considerable amount of frames to converge.
Revised: For example, (Gelada et al., 2019) showed how their dynamics and reward losses are inversely proportional to each other early on in the training, taking a considerable amount of frames to converge.

Original: We observe a similar pattern in Figure 8: the clustering loss first jumps up while the predictive loss is minimized, then the trend reverses, and both losses get minimized near the end of the training.
Revised: We observe a similar pattern in Figure 8: the clustering loss first jumps up while the predictive loss is minimized, then the trend reverses, and both losses get minimized near the end of the training.

Original: Case study: splitting other dynamics-aware losses Similar to the postulate of ATC (Stooke et 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.
Revised: Case study: splitting other dynamics-aware losses Similar to the postulate of ATC (Stooke et 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.

Original: To validate this, we conducted an additional set of experiments on two well-known self-supervised learning algorithms which leverage predictive information about future timesteps: Deep Reinforcement and InfoMax
Revised: To validate this, we conducted an additional set of experiments on two well-known self-supervised learning algorithms which leverage predictive information about future timesteps: Deep Reinforcement and InfoMax

Original: Learning (DRIML) (Mazoure et al., 2020) and Self-Predictive Representations (SPR) (Schwarzer et al., 2021).
Revised: Learning (DRIML) (Mazoure et al., 2020) and Self-Predictive Representations (SPR) (Schwarzer et al., 2021).

Original: We ran (i) the default version of the algorithms with joint RL and representation updates, as well as (ii) RL updates propagated only through the layers above the encoder.
Revised: We ran (i) the default version of the algorithms with joint RL and representation updates, as well as (ii) RL updates propagated only through the layers above the encoder.

Original: Qualitative assessment of clusters Figure 9 shows, for 5 environments, 4 randomly sampled states for 2 behavioral clusters (4 clusters for Starpilot – differences between clusters are easier to visualize in this environment).
Revised: Qualitative assessment of clusters Figure 9 shows, for 5 environments, 4 randomly sampled states for 2 behavioral clusters (4 clusters for Starpilot – differences between clusters are easier to visualize in this environment).

Original: Note that clustered states go beyond visual similarity, and capture action sequences, agent position, presence of enemies and even topological equivalence of various levels.
Revised: Note that clustered states go beyond visual similarity, and capture action sequences, agent position, presence of enemies and even topological equivalence of various levels.

Original: The choice of environments for this demonstration is dictated by the nature of the action space, e.g. projectiles in Starpilot and path tracing in Miner allow to better visualize agent’s behavior.
Revised: The choice of environments for this demonstration is dictated by the nature of the action space, e.g. projectiles in Starpilot and path tracing in Miner allow to better visualize agent’s behavior.

Original: 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).
Revised: 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).

Original: Figure 10 shows the t-SNE of from randomly sampled states along the CTRL training path on
Revised: Figure 10 shows the t-SNE of from randomly sampled states along the CTRL training path on

Original: Starpilot – embeddings learned by CTRL can be seen to concentrate into distinct clusters and around their respective centroids. 
Revised: Starpilot – embeddings learned by CTRL can be seen to concentrate into distinct clusters and around their respective centroids. [Table]

Original: 8. S HOWCASE : S LOW CLUSTERING CONVERGENCE LEADS TO BETTER GENERALIZATION
Revised: 8. S HOWCASE : S LOW CLUSTERING CONVERGENCE LEADS TO BETTER GENERALIZATION

Original: Trajectory clustering is key to representation learning not only in CTRL but also in a prior method,
Revised: Trajectory clustering is key to representation learning not only in CTRL but also in a prior method,

Original: Proto-RL.
Revised: Proto-RL.

Original: However, while Proto-RL uses it to pretrain a representation which it then keeps frozen during RL, CTRL applies clustering to evolve the representation as RL progresses.
Revised: However, while Proto-RL uses it to pretrain a representation which it then keeps frozen during RL, CTRL applies clustering to evolve the representation as RL progresses.

Original: This raises a question: how important is online clustering convergence rate for learning a good representation?
Revised: This raises a question: how important is online clustering convergence rate for learning a good representation?

Original: Intuitively, if 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[Table] induced by improved policies.
Revised: Intuitively, if 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.

Original: Therefore, it seems crucial to learn the behavioral similarities at a rate that allows cluster centroids to adapt to the value improvement path (Dabney et al., 2021).
Revised: Therefore, it seems crucial to learn the behavioral similarities at a rate that allows cluster centroids to adapt to the value improvement path (Dabney et al., 2021).

Original: To validate this hypothesis [Table] we conducted an analysis of the correlation between the clustering quality and thetest performance as a function of training progress on three representative games [Table]
Revised: To validate this hypothesis, we conducted an analysis of the correlation between the clustering quality and the

Original: 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.
Revised: 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.

Original: Results shown in Figure 11 provide evidence that picking the unsupervised learning procedure which converges the fastest (i [Table] e. uses the lowest temperature β ) does not necessarily lead to the best generalization performance.
Revised: Results shown in Figure 11 provide evidence that picking the unsupervised learning procedure which converges the fastest (i.e. uses the lowest temperature β ) does not necessarily lead to the best generalization performance.

Original: Based on Figure 11, we conjecture that fast clustering convergence hinders the performance of the RL agent due to clusters being fixed early on and not adapting to the distribution shift induced by the evolving RL policy.
Revised: Based on Figure 11, we conjecture that fast clustering convergence hinders the performance of the RL agent due to clusters being fixed early on and not adapting to the distribution shift induced by the evolving RL policy.

Original: 8. S HOWCASE : L EARNING BEHAVIORAL SIMILARITIES CAPTURES LOCAL PERCEPTUAL CHANGES
Revised: 8. S HOWCASE : L EARNING BEHAVIORAL SIMILARITIES CAPTURES LOCAL PERCEPTUAL CHANGES

Original: To demonstrate the importance of identifying behavioral similarities, we designed a toy example problem with 5 behavioral clusters, where clustering the behaviors correctly leads to finding a near-optimal policy.
Revised: To demonstrate the importance of identifying behavioral similarities, we designed a toy example problem with 5 behavioral clusters, where clustering the behaviors correctly leads to finding a near-optimal policy.

Original: Our example problem is based on the standard Ising model 2 – 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 a temperature parameter 1 /β .
Revised: Our example problem is based on the standard Ising model 3 – 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 a temperature parameter 1 /β .

Original: We randomly initialize 5 Ising models, each parametrized by an inverse temperature parameter on a uniform grid β ∈ [0 . 01 , 0 . 3] .
Revised: We randomly initialize 5 Ising models, each parametrized by an inverse temperature parameter on a uniform grid β ∈ [0 . 01 , 0 . 3] .

Original: The system state is given by the state of all 5 models, and all models evolve in parallel at every step.
Revised: The system state is given by the state of all 5 models, and all models evolve in parallel at every step.

Original: At each timestep, the agent needs to choose one of the models, and has 5 actions corresponding to these choices.
Revised: At each timestep, the agent needs to choose one of the models, and has 5 actions corresponding to these choices.

Original: At timestep t , the agent is allowed to observe only the state of the model it chose at this step, and gets a reward based only on this model’s state.
Revised: At timestep t , the agent is allowed to observe only the state of the model it chose at this step, and gets a reward based only on this model’s state.

Original: The reward yielded by Ising model i at timestep t is given by r t = −|| s i,t − G || 22 , where G is a goal state.
Revised: The reward yielded by Ising model i at timestep t is given by r t = −|| s i,t − G || 22 , where G is a goal state.

Original: For a given problem instance, we sample G randomly by instantiating a 6th Ising model with an unknown inverse temperature parameter β ∗ and letting G be its final configuration after evolving it for a random number of steps.
Revised: For a given problem instance, we sample G randomly by instantiating a 6th Ising model with an unknown inverse temperature parameter β ∗ and letting G be its final configuration after evolving it for a random number of steps.

Original: Figure 12 outlines the experimental setting for this study case.
Revised: Figure 12 outlines the experimental setting for this study case.

Original: The 5 behavioral clusters in our setting correspond to the 5 Ising models.
Revised: The 5 behavioral clusters in our setting correspond to the 5 Ising models.

Original: 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 β ∗ and 2) choose that model and collect the corresponding rewards.
Revised: 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 β ∗ and 2) choose that model and collect the corresponding rewards.

Original: Table 5 outlines the results we obtained by deploying CTRL with different number-of-clusters parameter value.
Revised: Table 5 outlines the results we obtained by deploying CTRL with different number-of-clusters parameter value.

Original: 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 score can be used to set the true number of clusters in CTRL which, in turn, corresponds to the highest-return policy discovered by CTRL. 
Revised: 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 score can be used to set the true number of clusters in CTRL which, in turn, corresponds to the highest-return policy discovered by CTRL. E = E = E =

Original: [Table]
Revised: E =...

Original:  =
Revised: E =

Revised: Returns -0.

Original: [Table]
Revised: [Table]

Original:  A DDITIONAL THEORETICAL FINDINGS
Revised: 8. A DDITIONAL THEORETICAL FINDINGS

Original: Do uncorrelated local changes to state embeddings affect the clustering?
Revised: Do uncorrelated local changes to state embeddings affect the clustering?

Original: Theorem 8.

Revised: Theorem 8.

Original: Let M be an MDP and let v ∈ V be a dynamics embeddings in M .
Revised: Let M be an MDP and let v ∈ V be a dynamics embeddings in M .

Original: Define(cid:40)
Revised: Define

Original: [Equation]
Revised: [Equation]

Original: (8) and pick δ 1 s.t.
Revised: and pick δ 1 s.t.

Original: it lies on the positive half-plane spanned by E ⊤ j − E ⊤ j ′ for some 1 ≤ j ′ ≤ E .
Revised: it lies on the positive half-plane spanned by E ⊤ j − E ⊤ j ′ for some 1 ≤ j ′ ≤ E .

Original: Then, v ′ =
Revised: Then, v ′ =

Original: v + δ and v belong to the same partition j .
Revised: v + δ and v belong to the same partition j .

Original: It becomes apparent from the above statement that perturbations to a single state or groups of state embeddings do not modify the partition membership as long as their direction aligns with that of the cluster embeddings.
Revised: It becomes apparent from the above statement that perturbations to a single state or groups of state embeddings do not modify the partition membership as long as their direction aligns with that of the cluster embeddings.

Original: 8. P ROOFS
Revised: 8. P ROOFS

Original: Throughout this section, we assume that the policy π is fixed, and that CTRL optimizes L clust only.
Revised: Throughout this section, we assume that the policy π is fixed, and that CTRL optimizes L clust only.

Original: Proof 1 (Theorem 8.1)
Revised: Proof 1 (Theorem 8.1)

Original: For two dynamics embeddings to be assigned to the same cluster j , the following should hold
Revised: For two dynamics embeddings to be assigned to the same cluster j , the following should hold

Original: [Equation]
Revised: [Equation]

Original: for any 1 ≤ j ′ ≤ E s.t.
Revised: for any 1 ≤ j ′ ≤ E s.t.

Original: j ′ ̸
Revised: j ′ ̸

Original: = j .
Revised: = j .

Original: [Equation]
Revised: [Equation]

Original: Taking the difference between both equations yields the necessary condition for two dynamics to belong to the same cluster
Revised: Taking the difference between both equations yields the necessary condition for two dynamics to belong to the same cluster

Original: [Equation]
Revised: [Equation]

Original: Corollary 8.1.1 Let v , v ′ be two dynamics embeddings, and define δ = v ′ − v .
Revised: Corollary 8.1.1 Let v , v ′ be two dynamics embeddings, and define δ = v ′ − v .

Original: If v belongs to cluster j and j = arg max 1 ≤ j ′ ≤ E E ⊤ j ′ δ , then v ′ also belongs to cluster j .
Revised: If v belongs to cluster j and j = arg max 1 ≤ j ′ ≤ E E ⊤ j ′ δ , then v ′ also belongs to cluster j .

Original: Perturbations are of the form (cid:80) |V| i =

Revised: Perturbations are of the form (cid:80) |V| i =

Original: ( E ⊤ ij − E ⊤ ij ′ )
Revised: ( E ⊤ ij − E ⊤ ij ′ )

Original: δ i .
Revised: δ i .

Original: If δ = 0 , then the cluster assignment doesn’t change.
Revised: If δ = 0 , then the cluster assignment doesn’t change.

Original: Let v be of size kh = |V| .
Revised: Let v be of size kh = |V| .

Original: Define, without loss of generality(cid:40) δ
Revised: Define, without loss of generality

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: and pick δ 1 s.t.
Revised: and pick δ 1 s.t.

Original: it lies on the positive half-plane spanned by E ⊤ ij − E ⊤ ij ′ .
Revised: it lies on the positive half-plane spanned by E ⊤ ij − E ⊤ ij ′ .

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: which concludes the proof.
Revised: which concludes the proof.

Original: Proof 2 (Theorem 1) Since the W 1 metric is defined between distribution functions, we use v = P [ v ] throughout the proof to denote the probability distribution over elements of the dynamics vector v .
Revised: Proof 2 (Theorem 1) Since the W 1 metric is defined between distribution functions, we use v = P [ v ] throughout the proof to denote the probability distribution over elements of the dynamics vector v .

Original: In practice, this amounts to re-normalizing the representation.
Revised: In practice, this amounts to re-normalizing the representation.

Original: For two dynamics to be assigned to the same cluster j , the following has to hold:[Equation] V [Equation]
Revised: For two dynamics to be assigned to the same cluster j , the following has to hold:

Original: ⊤

Original: [Equation]
Revised: [Equation]

Original: (14) for any 1 ≤ j ′ ≤ E s.t.
Revised: for any 1 ≤ j ′ ≤ E s.t.

Original: j ′ ̸
Revised: j ′ ̸

Original: = j .
Revised: = j .

Original: Then, adding both inequalities yields, for all 1 ≤ j ≤ E|V| (cid:88) i = v i E ⊤ ji + |V| (cid:88) i = v ′ i E ⊤ ji ≥ |V| (cid:88)
Revised: Then, adding both inequalities yields, for all 1 ≤ j ≤ E

Original: i = v i E ⊤ j ′ i + |V| (cid:88) i = v ′

Original: i E ⊤ j ′ i |V| (cid:88)

Original: i = v i

Original: E ⊤ ji + |V| (cid:88)

Original: i = v i E ⊤ j ′

Original: i ≥ |V| (cid:88) i = v ′

Original: i E ⊤ ji + |V| (cid:88)

Original: i = v ′ i E ⊤ j ′ i |V| (cid:88)

Original: i = v i ( E ⊤ ji − E ⊤ j ′ i ) ≥ |V| (cid:88)

Original: i = v ′

Original: i ( E ⊤ ji − E ⊤ j ′ i ) |V| (cid:88)

Original: i = v i ( E ⊤ ji − E ⊤ j ′ i ) ≥ |V| (cid:88)

Original: [Equation]
Revised: [Equation]

Original:  and the constraint of two vectors belonging to the same cluster j becomes
Revised: 15) and the constraint of two vectors belonging to the same cluster j becomes

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: Now, denote E ( j ) := E j .
Revised: Now, denote E ( j ) := E j .

Original: Our constraint satisfaction problem can be written as
Revised: Our constraint satisfaction problem can be written as

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: By comparing Eq.
Revised: By comparing Eq.

Original: 7 with Eq.
Revised: 7 with Eq.

Original: 17, we observe that in our case, µ is restricted to the set of vectors in R |V| .
Revised: 17, we observe that in our case, µ is restricted to the set of vectors in R |V| .

Original: Therefore, we pick µ ∈ Γ( E ) , where Γ( E ) = { ω ∈ V : ω = E
Revised: Therefore, we pick µ ∈ Γ( E ) , where Γ( E ) = { ω ∈ V : ω = E

Original: ( j, i ) − E ( j ′ , i ) , 0 ≤ ω i ≤ 1 , cos( v − v ′ , ω ) ∈ [0 , π ] | 1 ≤ i ≤ |V| , 1 ≤ j ′ ≤ E } .
Revised: ( j, i ) − E ( j ′ , i ) , 0 ≤ ω i ≤ 1 , cos( v − v ′ , ω ) ∈ [0 , π ] | 1 ≤ i ≤ |V| , 1 ≤ j ′ ≤ E } .

Original: The set Γ( E ) is non-empty if max l,l ′
Revised: The set Γ( E ) is non-empty if max l,l ′

Original: ||
Revised: ||

Original: E l − E l ′ || ∞ ≤ 1 , which holds due to ℓ p norm ordering and since E is normalized in the ˜ Q scores expression.
Revised: E l − E l ′ || ∞ ≤ 1 , which holds due to ℓ p norm ordering and since E is normalized in the ˜ Q scores expression.

Original: Adopting this notation simplifies the previous expression to
Revised: Adopting this notation simplifies the previous expression to

Original: [Equation]
Revised: [Equation]

Original: Once again, recall that E is normalized.
Revised: Once again, recall that E is normalized.

Original: Therefore, we have
Revised: Therefore, we have

Original: [Equation]
Revised: [Equation]

Original:  which equivalently can be re -stated as (for e i ≥ e j WLOG):
Revised: 19) which equivalently can be re-stated as (for e i ≥ e j WLOG):

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: where we take, as an example, d ( i, i ′ )
Revised: where we take, as an example, d ( i, i ′ )

Original: =
Revised: =

Original: ||
Revised: ||

Original: e j || 2 (cid:0) E ij − E ij ′ (cid:1) .
Revised: e j || 2 (cid:0) E ij − E ij ′ (cid:1) .

Original: The final expression for the sufficient condition for two dynamics embeddings to belong to the same partition is
Revised: The final expression for the sufficient condition for two dynamics embeddings to belong to the same partition is

Original: [Equation]
Revised: [Equation]

Original: for d ( i, i ′ )
Revised: for d ( i, i ′ )

Original: =
Revised: =

Original: ||
Revised: ||

Original: e j || 2 (cid:0) E ij − E ij ′ (cid:1) , which is similar to the Wasserstein-1 distance under d , i.e.
Revised: e j || 2 (cid:0) E ij − E ij ′ (cid:1) , which is similar to the Wasserstein-1 distance under d , i.e.

Original: W d 1 ( v , v ′ ) .
Revised: W d 1 ( v , v ′ ) .

Original: We constructed an operator similar to F ( d ) in Ferns et al.
Revised: We constructed an operator similar to F ( d ) in Ferns et al.

Original: d can be computed by recursively applying F ( d ) at each v , v ′ ∈
Revised: d can be computed by recursively applying F ( d ) at each v , v ′ ∈

Original: V pointwise, which is similar to what is done in CTRL.
Revised: V pointwise, which is similar to what is done in CTRL.

Original: This concludes our proof and shows how our clustering procedure can be viewed as finding reward-free bisimulations.
Revised: This concludes our proof and shows how our clustering procedure can be viewed as finding reward-free bisimulations.

Original: However, note that the exact interpretation of reward-free bisimulation relation depends on how v is defined.
Revised: However, note that the exact interpretation of reward-free bisimulation relation depends on how v is defined.

Original: Taking v to 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 a different set of properties.
Revised: Taking v to 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 a different set of properties.