Cross-Trajectory Representation Learning for Zero-Shot Generalization in RL

Algorithm overview. For each trajectory batch, CTRL performs 4 operations:

1. 1.

   Apply the observation history (belief state) encoder $\phi$ to generate a low-dimensional reward-agnostic representation (“view”) of each trajectory.

2. 2.

   Group trajectories’ views into $C$ sets ($C$ is a tunable hyperparameter) using an online clustering algorithm with loss $\mathcal{L}_{\text{clust}}$ (Equation 4, Section 4.2).

3. 3.

   Using trajectory pairs selected from neighboring clusters, apply a predictive loss $\mathcal{L}_{\text{pred}}$ (Equation 6, Section 4.2) to encourage $\phi$ to capture cross-cluster behavioral similarities.

4. 4.

   Update $\phi$ with gradients of the total loss: $\mathcal{L}_{\text{CTRL}}=\mathcal{L}_{\text{clust}}+\mathcal{L}_{\text{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. We conduct ablations to show the effect of removing the clustering and predictive objectives of CTRL, with details in Appendix 8.3.

Clustering. $C$ clusters can be viewed as corresponding to $C$ latent “situations” in which an RL agent may find itself. Each situation is essentially a group of belief states. CTRL’s implicit hypothesis is that a given policy should behave roughly similarly across all belief states corresponding to the same “situation”, i.e., generalize across similar belief states. Under this hypothesis, an agent’s policy can be expected to produce $C$ sets of roughly similar trajectories. CTRL’s clustering step (#2 above) is an attempt to recover these trajectory sets. Since each trajectory consists of belief states, the purpose of $\mathcal{L}_{\text{clust}}$ is to force the encoder $\phi$ to compute belief state representations that make trajectories within each cluster look similar in the latent space.

Since we would like to evolve clustering online as new trajectory batches arrive, we employ a common online clustering algorithm, the Sinkhorn-Knopp procedure [Caron et al., 2020], which has been used as an auxiliary RL loss [e.g., Proto-RL, Yarats et al., 2021].

Cross-cluster prediction. Note, however, that the clustering loss emphasizes the recognition of behavioral similarities within clusters.

This may hurt generalization, as the resulting centroids may not faithfully represent behaviors encountered at test time. Our hypothesis is that encouraging encoder $\phi$ to induce latent-space similarities between trajectories from different but adjacent clusters will increase its ability to recognize behaviors in unseen test trajectories.

While there are several ways to encourage cross-cluster representational similarity, using a mechanism similar to MYOW [Azabou et al., 2021] on trajectories drawn from neighboring clusters captures this idea particularly well. Namely, to get the cross-predictive loss $\mathcal{L}_{\text{pred}}$, we sample trajectory view pairs from neighboring clusters and apply the cosine-similarity loss to those pairs.

Using reward guidance without reward signal for representation learning. CTRL trains encoder $\phi$ only using the gradients of $\mathcal{L}_{\text{CTRL}}$; the RL algorithm’s loss $\mathcal{L}_{\text{RL}}$ trains only the policy head. Thus, encoder $\phi$ is isolated from the previously observed dangers of overfitting to the reward function [Raileanu and Fergus, 2021] that shapes $\mathcal{L}_{\text{RL}}$. 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 $\mathcal{L}_{\text{RL}}$’s gradients, which are reward-dependent.

Below we describe the details of each of CTRL’s steps, with CTRL’s pseudocode presented in Algorithm 1 in Appendix 8.1. While in general the agent’s belief state at step $t$ of a trajectory is the entire observation history $h_{t}=(o_{1},\ldots,o_{t})$, in the rest of the section we will assume $h_{t}=(o_{t})$ and, in a slight abuse of notation, use $\phi(o_{t})$ instead of $\phi(h_{t})$ to simplify explanations²²2While, in theory, the Procgen suite is indeed a POMDP, most RL algorithms take the most recent observation as the belief state – a simplification which was shown not to hinder ZSG on Procgen [Cobbe et al., 2020].. We emphasize, however, that CTRL equally applies in settings where the agent uses a much longer history of observations as its state. In this case, $\phi$ would be recurrent or process stacks of frames.

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

Generating low-dimensional trajectory views with encoder $\phi$. CTRL’s input in each epoch is a trajectory batch $\{traj_{i}\}_{i=1}^{B}$ of size $B$. Assume all trajectories in the batch have the same length $L$. For an integer hyperparameter $T\leq L$, for each trajectory $traj_{i}=(o_{0},a_{0},r_{0},\ldots o_{T_{b}},a_{L},r_{L})$ we independently and uniformly sample a subset of its steps $\tau_{i}=t_{1},\ldots,t_{T}$ to form a subtrajectory $(o_{t_{1}},a_{t_{1}},r_{t_{1}},\ldots o_{t_{T}},a_{t_{T}},r_{t_{T}})$. We then encode this subtrajectory as

where $FiLM(\phi(o_{t_{j}}),a_{t_{j}})$ is a common way of combining representations of different objects, akin to conditioning Perez et al. [2018], and the resulting vector $\bm{u}_{i}^{(\tau_{i})}$ is in a low-dimensional space $\mathcal{U}$. Note two aspects of the process of generating these vectors: (1) it drops rewards from the original trajectory and (2) it critically relies on $\phi$ whose parameter values are learned from previous epochs. Vectors $\bm{u}_{i}^{(\tau_{i})}$ produced in this way are the trajectory views that the next steps of CTRL operate on.

Clustering trajectory views. CTRL groups trajectories from the epoch’s batch by clustering the set of their views $\{\bm{u}_{i}^{(\tau_{i})}\}_{i=1}^{B}$. 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]. Since CTRL operates online, Sinkhorn-Knopp is better-suited for the task than other clustering methods.

The clustering branch computes two views of each input $\bm{u}_{i}^{(\tau_{i})}$ in a cascading fashion: first by passing it through a clustering encoder $\psi_{\text{clust}}\mathrel{\mathop{\mathchar 58\relax}}\mathcal{U}\to\mathcal{V}$, e.g. an RNN, to obtain a lower-dimensional view $\bm{v}_{i}=\psi_{\text{clust}}(\bm{u}_{i}^{(\tau_{i})})$, and then by passing $\bm{v}_{i}$ through yet another network, an MLP $\theta_{\text{clust}}\mathrel{\mathop{\mathchar 58\relax}}\mathcal{V}\to\mathcal{W}$, to produce view $\bm{w}_{i}=\theta_{\text{clust}}(\bm{v}_{i})$. The parameters of $\psi_{\text{clust}}$ and $\theta_{\text{clust}}$ are learned through the epochs jointly with $\phi$’s. Like $\bm{u}_{i}^{(\tau_{i})}$, each $\bm{v}_{i}$ and $\bm{w}_{i}$ is a view of trajectory $i$; the approach then consists in projecting $\bm{v}_{i}$’s and $\bm{w}_{i}$’s onto centroids of $C$ clusters in two different ways and then computes a clustering loss that enforces consistency between $\bm{v}_{i}$’s and $\bm{w}_{i}$’s cluster projections.

Specifically, we represent the centroid of each cluster $c$ with a vector $\bm{e}_{c}\in\mathcal{V}$, which are stacked into a matrix $\mathbf{E}$. These vectors are additional parameters in the joint optimization problem CTRL solves. They can be regarded as views of $C$ typical behaviors around which the trajectories’ views are regrouped. To project trajectory $i$’s $\bm{v}_{i}$ views onto behavioral centroids learned from previous trajectory batches, CTRL computes a vector of soft assignments of $\bm{v}_{i}$ to each centroid $\bm{e}_{c}$:

and forms a $B\times C$ matrix $\mathbf{Q}$ whose $i$-th row is the soft assignment of $\bm{v}_{i}$. The resulting assignments may be very unbalanced, with most probability mass assigned to only a few clusters. Applying the Sinkhorn-Knopp algorithm solves this issue by iteratively re-normalizing $\mathbf{Q}$ in order to obtain a more equal cluster membership [Cuturi, 2013], where the degree of re-normalization is controlled by a temperature parameter $\beta$. The output of this operation is a matrix $\tilde{\mathbf{Q}}$.

For each view $\bm{w}_{i}$ (the projection of trajectory view $\bm{v}_{i}$) CTRL computes the logarithm of its soft cluster assignments and treats these vectors as rows of another $B\times C$ matrix $\mathbf{P}$:

Finally, we compute the cross entropy between $\tilde{\mathbf{Q}}$ and $\mathbf{P}$, which measures their inconsistency. This measure is taken as the clustering loss:

Encouraging cross-cluster behavioral similarity. Note, however, that the clustering loss in the above step emphasizes the recognition of behavioral similarities within clusters. This may hurt generalization, as the resulting centroids may not faithfully represent behaviors encountered at test time. Our hypothesis is that encouraging encoder $\phi$ to induce latent-space similarities between trajectories from different but adjacent clusters will increase its ability to recognize behaviors in unseen test trajectories.

While there are several ways to encourage cross-cluster representational similarity, using a mechanism similar to MYOW [Azabou et al., 2021] on trajectories drawn from neighboring clusters captures this idea particularly well. Namely, to get the cross-predictive loss $\mathcal{L}_{\text{pred}}$, we sample trajectory view pairs from neighboring clusters and apply the cosine-similarity loss to those pairs.

To implement this idea, we define a measure of cluster proximity via a matrix $\mathbf{D}$ of cosine similarities between cluster centroids: for clusters $k$ and $l$, $\mathbf{D}_{kl}=||\bm{e}_{k}-\bm{e}_{l}||^{2}_{2}$  [Grill et al., 2020]. Recall that in the previous step, the clustering branch computed a matrix $\mathbf{Q}$ whose rows $\mathbf{Q}_{i}$ are soft assignments of trajectory $i$’s view $\bm{v}_{i}$ to clusters. In this step, we convert these soft assignments to hard ones by associating a trajectory’s view $\bm{v}_{i}$ with cluster $c_{i}=\text{argmax}_{1\leq c^{\prime}\leq C}\tilde{\mathbf{Q}}_{i}$ and treating a cluster $c$ as consisting of trajectories with indices in the set $\mathbb{T}_{c}=\{i\mid c=\text{argmax}_{1\leq c^{\prime}\leq C}\tilde{\mathbf{Q}}_{i}\}$. To assess how predictive a trajectory embedding $\bm{u}_{i}$ is of a trajectory embedding $\bm{u}_{j}$, like in the clustering step we will use two special helper maps, $\psi_{\text{pred}}\mathrel{\mathop{\mathchar 58\relax}}\mathcal{U}\to\mathcal{V}$ to obtain a reduced-dimensionality view $\bm{v}^{\prime}=\psi_{\text{pred}}(\bm{u})$ and $\theta_{\text{pred}}$ to further project $\bm{v}^{\prime}$ to $\bm{w}^{\prime}=\theta_{\text{pred}}(\bm{v}^{\prime})$.

CTRL proceeds by repeatedly sampling trajectories, which we call anchor trajectories, from the batch, with their associated embeddings $\bm{u}$. For each anchor trajectory $n$, consider $K$ clusters $c_{1},\ldots,c_{K}$ nearest to $n$’s cluster $c_{n}$, as defined by the indices of $K$ largest values in row $c_{n}$ of matrix $\mathbf{D}$ (we exclude $c_{n}$ itself when determining $c_{n}$’s nearest clusters). Borrowing ideas from the MYOW approach Azabou et al. [2021], CTRL mines a view for $\bm{u}_{n}$ by randomly choosing a trajectory with embedding $\bm{u}^{(n)}_{c_{k}}$ from each of the neighboring clusters and computing its view $\bm{v}^{\prime}_{c_{k}}=\psi_{\text{pred}}(\bm{u}^{(n)}_{c_{k}})$. We call the neighbors’ views $\bm{v}^{\prime}_{c_{1}},\ldots,\bm{v}^{\prime}_{c_{K}}$ trajectory $n$’s mined views.

For the final operation in this step, CTRL computes trajectory $n$’s predictive view $\bm{w}^{\prime}_{n}=\theta_{\text{pred}}(\psi_{\text{pred}}(\bm{u}_{n}))$ and measures the distance from it to trajectory $n$’s mined views:

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

Updating encoder $\phi$ using reward guidance without reward signal for representation learning. Note that CTRL’s total loss $\mathcal{L}_{\text{CTRL}}$ depends on the parameters of encoder $\phi$ as well as of clustering networks $\phi_{\text{clust}}$ and $\theta_{\text{clust}}$, prediction networks $\phi_{\text{clust}}$ and $\theta_{\text{clust}}$, and cluster centroids $\bm{e}_{c}$, $1\leq c\leq C$. In each epoch, CTRL updates all these parameters to minimize $\mathcal{L}_{\text{CTRL}}$.

We implemented all algorithms on top of the IMPALA architecture, which was shown to perform well on Procgen [Cobbe et al., 2020]. 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). Similarly, DIAYN was also run with the pre-training phase of 8M and then RL only objective for the second 8M phase.

Like Proto-RL and DIAYN, PSE needed extra training budget and extra adjustments for a fair comparison to CTRL (see Section 2). Agarwal et al. [2021] used PSE only with Soft Actor-Critic, which doesn’t perform well on Procgen. Therefore, we carried over PSE’ available implementation 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.

PSE assumes being given policies for training problem instances (Procgen levels). Like Agarwal et al. [2021], we ran PSE using both random and high-quality pretrained policies for these problems. In the latter case, we pretrained expert policies for the first 40 levels of Procgen to generate training trajectories for PSE. Each level’s expert was trained on 0.5M environment steps. Note that this is much less than 8M steps 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. We did this only for the first 40 levels because even for 40 levels this took 20M training steps per game and had to be done for 16 games. We don’t believe more than 40 experts per game would have made a difference.

Given policies for the training levels, PPO+PSE’s training on Procgen mimicked CTRL’s: PPO+PSE trained by interacting with the first 200 levels for 8M steps and was evaluated on the rest. However, during the training, PPO+PSE sampled an additional 1M interactions from the pretrained traiing-level policies. the process was repeated for all 16 games, for 10 seeds each.

We did hyperparameter grid search on PPO+PSE’s hyperparameters for PSE – loss coefficient values of $(0.1,1,2)$ and temperature $(0.1,0.3,0.7)$. PPO’s hyperparameters were the same as in CTRL.

Thus, due to the need to pretrain and gather data with per-level expert policies, PPO+PSE received $0.5\cdot 20M+1M=21M$ extra environment interactions compared to CTRL, i.e., used $21M/8M=2.6\times$ more training data than the latter.

Trajectory clustering is key to representation learning not only in CTRL but also in a prior method, Proto-RL. 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. This raises a question: how important is online clustering convergence rate for learning a good representation? 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. 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]. To validate this hypothesis, we conducted an analysis of the correlation between the clustering quality and the test performance as a function of training progress on three representative games.

Results shown in Figure 11 provide evidence that picking the unsupervised learning procedure which converges the fastest (i.e. uses the lowest temperature $\beta$) does not necessarily lead to the best generalization performance. 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.

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.

Our example problem is based on the standard Ising model³³3https://en.wikipedia.org/wiki/Ising_model – a $32\times 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/\beta$. We randomly initialize $5$ Ising models, each parametrized by an inverse temperature parameter on a uniform grid $\beta\in[0.01,0.3]$. The system state is given by the state of all 5 models, and all models evolve in parallel at every step. At each timestep, the agent needs to choose one of the models, and has 5 actions corresponding to these choices. At timestep $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. The reward yielded by Ising model $i$ at timestep $t$ is given by $r_{t}=-||s_{i,t}-G||_{2}^{2}$, where $G$ is a goal state. For a given problem instance, we sample $G$ randomly by instantiating a 6th Ising model with an unknown inverse temperature parameter $\beta^{*}$ and letting $G$ be its final configuration after evolving it for a random number of steps. Figure 12 outlines the experimental setting for this study case.

Table 5 outlines the results we obtained by deploying CTRL with different number-of-clusters parameter value. One can see that the largest improvement in silhouette score occurs from $E=4$ 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.

Throughout this section, we assume that the policy $\pi$ is fixed, and that CTRL optimizes $\mathcal{L}_{clust}$ only.