Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes an extension of the MYOW (Azabou et al 2021) self-supervised learning technique, combining it with SwAV (Caron et al 2021)’s method to perform online clustering, and adapt it for RL to assess generalisation performance on Procgen.
It compares against several recent baselines (DBC, PSE, CURL, Proto-RL and DIAYN), and outperforms them consistently on the tasks assessed (albeit by a small margin, depending on the reward scale of Procgen which I’m not extremely familiar with).
main_review: I found the paper interesting, but rather hard to understand what it really contributed and how several decisions came about. It is really only after I read MYOW and SwAV that I could understand several model architecture and loss choices, as the main text doesn’t cover them well enough and the Appendix is quite lacking. It also appeared to me that the proposed model is really close to MYOW and that fact wasn’t made as clear as I would have liked.


Additional comments and questions:

1. The introduction and background are very clear and set up the problem well. The literature review is thorough and recent, and the overview of the algorithm in section 4 is also clear.
2. My main issue was in understanding several model choices, when looking at Figure 1 especially. After looking at MYOW and SwAV, one can see how the overall model is basically MYOW with a different sequence of networks (a simpler one, which appears more plausible than MYOW to me, although see the next point) combined cleverly with SwAV’s soft-clustering model. CTRL also removes all multi-views aspects of these works, instead operating on trajectories directly.
   1. It would have been much clearer to spell out these differences and what you did explicitly, given I would still consider that CTRL is significantly different from either of these two previous works.
   2. One particular point which was confusing was the decision to make latent space be similar to adjacent clusters. This is taken directly from MYOW, but wasn’t clearly expressed here. It also appears slightly odd as there is no enforcement of a “global coverage” across the clusters (i.e. there is no extra repulsion term to far-away clusters), which to me would indicate that this method will have a chance to collapse (although one can be lucky by using target networks, ala BYOL).
   3. Similarly, the choice to have a cascade of networks creating the two embedding for the clustering is rather surprising (and I would expect it to collapse to degenerate solutions given you’re not applying them in swapped ways like in SwAV).
   4. The main text talks about trajectories, but T=2 if the Appendix is to be trusted. Hence in practice you are only clustering transitions. It would have been easier to present it that way (and the connection to bisimulation metrics would have been clearer).
3. Many details aren’t in the main text, and the Appendix is equally light in details:
   1. What is your encoder? Does it receive a history / have a recurrent state? Appendix 8.2 presents it as if it is instantaneous, which is quite a limitation.
   2. How is FiLM used exactly? There are infinitely many ways to add it in.
   3. It is very hard to follow the exact sequence of v/w/v’/w’ and how they are computed. One needs to carefully read SwAV to build an intuition, but this shouldn’t be that hard. No details are provided on what these are, apart from “\psi_clust is a RNN” and “\theta_clust is a MLP”
   4. MYOW explicitly used predictors everywhere and were never regressing towards the latent used to cluster (i.e. `y = f_theta(x)` in their paper. See Figure 1, the augmented and clustering loss are applied on z and v, respectively.). Instead, you directly force `w’` to be predictable of `v’_c_i`. It would be valuable to discuss why you believe this is more appropriate (MYOW’s argue that this worked worst for them).
   5. It would be quite valuable to bring back some of Section 8.2 into the main text, and add many more details to clarify what the model really is, I don’t believe this is reproducible in this current state.
4. Table 1 is great and I would again like to flag that the number of baselines available are clearly to a great standard. This must have been hard to do well, so congratulations!
   1. One caveat is that you are only reporting Evaluation performance on held out tasks (which is good, that was your metric), but I always would have appreciated seeing the performance on the training distribution. This would make it clearer if the effect is about improved generalization, or about the policies themselves being better or worse for particular baselines/ablations.
5. Section 6.2 is quite anecdotal, and it isn’t clear that the effect discussed is that significant?
6. Section 6.3 feels out of place and is quite underwhelming after having covered the results on Procgen. 
   1. It is a very peculiar toy situation, which isn’t that standard, and I’m not entirely sure what it’s trying to demonstrate in its current state. 
   2. Can you use the same silhouette score change to tune the number of clusters in Procgen too? Given it’s a post-hoc analysis, this doesn’t really appear like a real solution to this problem.
7. Figure 4 in the Appendix has T=3 and T=5, but they show no difference? There is also no discussion of Figure 4, the section in 8.3 appears empty.
8. As a small point, you keep referring to “views”, but this is overloaded and isn’t as appropriate as you do not actually have augmentations like in both previous work? Once again makes it hard to follow but should be easy to fix.
9. You might want to add a citation to Van der Pol et al 2020 (https://arxiv.org/abs/2002.11963) for the bisimulation work, as they had a slightly different take than Zhang et al 2021.

summary_of_the_review: Overall, I think this is nice paper, with good baselines and clear results, however, its presentation right now is quite lacking and hence I feel it may need some work before being ready for publication.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: This paper addresses the zero-shot generalization (ZSG) problem in reinforcement learning (RL). Specifically, it focuses on learning state representations that can generalize well. To this end, the authors propose a self-supervised learning method, cross trajectory representation learning (CTRL), that captures behaviorally similarity in the trajectory representations. The CTRL objective has two components, one for online clustering and the other one for cross-cluster prediction. CTRL is shown to be connected to bisimulation metrics in RL. The empirical study in the ProcGen benchmark shows that CTRL performs better than existing methods in the literature. Further experiments demonstrate the importance of adaptive online clustering and shed light on how to choose the number of clusters as a hyperparameter.
main_review: **Originality and significance**:
This paper addresses zero-shot generalization, which is an important problem in RL. The idea of using unsupervised representation learning techniques to avoid overfitting in RL is not new. But the proposed method, CTRL, is somehow novel in that it uses both online clustering and cross-cluster prediction for representation learning, which has not been explored in the literature.

**Quality**:
The overall quality of the paper is below the threshold for acceptance. Below are detailed comments.
1. The setup for the ProcGen experiment presented in Section 6.1 is unclear. ProcGen has two difficulty modes and the main text is not clear about which one was used in the experiment. I guess the easy mode was used based on the number of training levels. Moreover, the 8 million steps training budget is unusual. The original ProcGen paper recommended 25 million steps for the easy mode. The choice of reporting test performance after 8 million steps of training seems arbitrary and needs further justification.
2. I have some questions regarding the results in Table 1. Although the main text says "...DAAC also exhibits good generalization performance...", DAAC in fact performs worse than PPO (6 wins, 9 loses, and 1 tie in 16 games). These results conflict with the results reported in the original paper by Raileanu and Fergus. This is a bit surprising as Raileanu and Fergus opensourced their code. One possible explanation is that the discrepancy is due to the lower training budget, but it then leads back to the first question of why using 8 million steps rather than the recommended 25 million steps. Overall, the discrepancy between the reported results and the published results makes Table 1 less convincing.
3. The experiment in Section 6.1 is very helpful. It verifies the importance of adapting the clusters online as the policy improves.
4. I have some difficulty understanding the purpose of Section 6.3. The main text says "To demonstrate the importance of identifying behavioral similarities..." But after reading this section through, all I got was a way to select the hyperparameter for the number of clusters. I will be grateful if the authors can provide some clarification.

**Clarity**:
The overall flow of the paper is easy to follow. But there are critical details missing in the method description which makes it hard to understand. Although the Sinkhorn-Knopp procedure and MYOW are existing methods, the authors should provide minimum details in the *main text* so that the readers can understand what is going on. For example, the concrete equations for computing $L_{\text{clust}}$ and $L_{\text{pred}}$ should be presented. Figure 1 can be explained in more detail. Currently, it contains unexplained notations like $\textbf{e}_{c}$.

**Minor issues & typos**:
* The references seem truncated. For example, all the citations to Zhang _et al_ are not shown in the references.
* The citation to Kendall _et al_ in the 3rd line of Section 1 should use \citep.
* The 7th line on page 3: "bisimation" --> "bisimulation".
* The 4th last line on page 5: the citation to Raileanu and Fergus should use \citep.
* Bullet point 2 in Definition 1: if $c$ is a state, what does $s' \in c$ mean? Is it a typo?
* The line below Definition 1: the citation to Ferns _et al_ should use \citet. 
summary_of_the_review: I think this paper is mainly held back by two factors. First, critical details are missing in the main text which makes the proposed method hard to understand. Second, I have some questions regarding the main empirical results that may undermine the empirical significance of this work. Therefore, I don't think this paper is ready to publish yet.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes a new zero-shot generalization method for RL that uses a novel clustering+prediction loss to learn a cross-trajectoriy encoder. The learned encoder is supposed to map behaviorally similar observations to similar representations. The algorithm generally outperforms representation RL baselines in experiment.
main_review: Strengths:

1. This paper is well-written and easy to follow.
2. The proposed method of training the encoder using clustering and prediction losses seems novel to me.
3. Experimental results in all 16 tasks in the Procgen benchmark show that the proposed CTRL generally achieve zero-shot generalization.
4. The additional analysis on clustering convergence and local perceptual changes are helpful for practical usage of CTRL.

Weaknesses:
1. A lot of details are omitted from the main paper, e.g. the detailed definitions of the clustering loss and the prediction loss.
2. The proposed method does not significantly outperform baselines in many tasks. Other than bigfish, bossfight and starpilot, the improvement made by CTRL compared to baselines is relatively marginal.

Other questions:
1. Is it possible that the representation loss leads to a degenerate solution? For example, the encoder maps all observations to the same representation. I do not find a constraint to avoid degeneration.
2. What is the relation between this proposed clustering+prediction loss and other contrastive learning based methods such as CURL? 
3. How does it compare to the Darla method[1]?

[1] Higgins, Irina, et al. Darla: Improving zero-shot transfer in reinforcement learning.

summary_of_the_review: This paper proposes a novel idea of learning representation with a clustering+prediction loss, which makes the representation generalize to novel observations from the same task distribution. The experimental results are convincing, although a little marginal in many environments. The proposed method has some potential to help representation learning in RL and zero-shot generalization.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper addresses the problem of zero short generalization (ZSG) in RL by improving the generalization of the encoder. A new self-supervised learning objective is proposed to encourage the encoder to map behaviorally similar observations to similar representations without the use of reward signals. The overall model is named Cross Trajectory Representation Learning (CTRL), which achieves better generalization performance on the challenging Procgen benchmark suite.
main_review: #### **Strengths**
- The paper tackles a very important problem in the deep RL community.
- While various components that constitute CTRL were previously introduced in other works, I believe that the paper presents a novel combination of these methods and applies them in a novel problem setting.
- The proposed method is simple and can be easily plugged in to solve ZSG problems other than Procgen.
- The experiments are extensive, as the authors have compared CTRL to various baselines, including bisimulation metrics and SSL-based methods. The proposed method also outperforms quite significantly these baselines in most tasks.
- The paper is well written and easy to understand.

#### **Weaknesses**
- The major issue with the paper is the lack of discussions and explanations on why the proposed objective functions allow the encoder to recognize the behavioral similarity. Consider a scenario where a robot arm needs to pick up a particular object on a table. Each time the robot performs its task, the surface of the table will randomly change its color, which means the distribution of tasks is now the distribution of the colors. In this imaginary scenario, maximizing $\mathcal{L}_{\text{clust}} + \mathcal{L}_{\text{pred}}$ might learn an encoder that only encodes the color information, which is not what we want to achieve. If we utilize the reward information, however, we can avoid this behavior by forcing the encoder to not encode what is not useful to predict the reward. So the reward information in this case is actually helpful for learning a good representation.
- The visualization in Figure 8 in the Appendix shows that in practice CTRL does recognize the behavioral similarity. But again, I want to understand more on why this is the case, and why the encoder does not 'cheat' in maximizing the proposed objective function.

#### **Minor comments**
- In the first sentence of the second paragraph on Page 4, the distribution should be $P(\mathcal{O}, \mathcal{Z}, \mathcal{S})$. Same problem for the following sentence.
- 'repeatedly' was repeated twice in the caption for Figure 1.

summary_of_the_review: In general, the paper is interesting. The proposed method is novel and while being simple, it works effectively and outperforms the baselines. However, I want the authors to discuss in more detail about why their method works well in practice while in theory it can fail.
review_title: null

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

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

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

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