Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: The authors propose an interesting direction toward generalization in reinforcement learning. The authors propose a policy similarity metric (PSM) for measuring behavioral similarity between states. States are considered similar if the optimal policies in those states (as well as future states) are similar. The authors present a contrastive representation procedure to embed state similarity metrics, which are instantiated with PSM to yield policy similarity embeddings (PSEs). 

I like the idea of looking at the similarity between policies by examining the difference in local and long-term behavior. Furthermore, I like that the contrastive metric embedding approach is relatively simple and that PSEs can be learned alongside RL using an auxiliary objective.

Questions and comments:
1. What informed the train test split in the problem setup on page 4?
2. I would like to see how this approach performs on environments specifically designed to test generalization, such as the ProcGen [1] and MineRL [2] environments.

[1] Cobbe, et al. Leveraging procedural generation to benchmark reinforcement learning. arXiv:1912.01588, 2019.
[2] Guss, et al. The MineRL competition on sample efficient reinforcement learning using human priors. arXiv:1904.10079, 2019.
review_title: Official Blind Review #4

final_score: rating: 7: Good paper, accept
review_content: Summary: 

This work proposes a policy similarity metric (PSM) which differs slightly from the \pi-bisimulation metric: while the \pi-bisimulation metric defines the distance between two states in terms of the difference between the expected rewards obtained by following policy \pi, the proposed policy similarity metric measures distance based on the similarity between the optimal policies at each state. The policy similarity metric can then be used to learn policy similarity embeddings (PSEs) via contrastive representation learning. These embeddings should in principle improve generalization in RL. 



Strengths: 

- The proposed policy similarity metric is well motivated for generalization in RL, which is an important problem. I also appreciate the theoretical upper bound on the suboptimality of transferring policies to a new environment.

- The experiments are thorough and explore very diverse tasks. They highlight how PSEs can be used in imitation learning, LQR under spurious correlations, and RL settings; and can be combined with existing data augmentation techniques, including RandConv and DrQ. 

- The ablation studies and visualizations demonstrate, both quantitatively and qualitatively, the individual importance of the policy similarity metric and the contrastive metric embeddings.

- The paper is very clearly written and well structured.



Weaknesses:

- The main weakness of the proposed policy similarity metric is that its computation requires optimal trajectories or approximately optimal trajectories (obtained, for example, through RL) from at least two training MDPs.

- Building upon the previous point, I am curious about the utility of PSEs learned from increasingly suboptimal policies. I think including such a study would greatly help understand the sensitivity of learning PSEs to quality of the policies.

- In settings like the jumping task from pixels, because of the reliance of the PSM on positive examples, data augmentation is required for good performance. 



Recommendation:

I am recommending to accept this paper. I think the work is well motivated, approaches an important problem, and conducts very thorough experiments. I would still like to understand how useful PSEs are when learned from increasingly suboptimal policies.



Questions:

- For the Distracting Control Suite tasks, are the “DrQ + PSEs” policies trained with both DrQ and PSEs or only PSEs? Specifically, I’m trying to understand if the gains are only from PSEs as opposed to the combination of PSEs and data augmentation.



Updates:

Thank you for the clarifications, new study, and theoretical justification. The new results demonstrate the utility of PSEs even when policies are not exactly optimal. My concerns were thoroughly addressed by the authors’ response, and I’ve updated my score accordingly.
review_title: Well-motivated representation learning approach for RL generalization

final_score: rating: 6: Marginally above acceptance threshold
review_content: This paper defines a new metric for calculating similarities between states. The policy similarity metric is basically based on the $\pi$-bisimulation metrics, but the l1 distance term of rewards is modified to the distance between states transitioned by the grounded policy. The experiments performed with several benchmarks, and the gain seems to be consistent.

[Quality]

The paper is clearly written overall and well arranged.

[Originality & Significance]

The originality and significance are sufficient for acceptance, as the authors devised a novel metric.

[Strengths]
+ Even though I didn’t thoroughly check all of the derivations, the metric seems to be more precise than $\pi$-bisimulation, similarity between states originated from different environments.
+ This approach can be easily applied to various RL models.

[Weaknesses]
- In Table 1, experiments with RandConv only contain RandConv and RandConv+PSEs. It should be better to compare to other methods, at least Bisimulation, to show how much the PSE model performs well. Not only Table 1, I feel overall comparison results are not sufficient.
- The Distracting Control Suite dataset, which is the core dataset in the experiments, is not currently available, so verifying the results with this dataset is difficult.

[Questions or Additional Comments]
1. Figure 2 is not self-explanatory and hard to understand. I think adding axes or an additional explanation in the caption is required.
2. Please elaborate on the reason for using the Distracting DM Control Suite dataset rather than other datasets.
3. According to Tables G.3 and G.4, hyperparameters are quite different in each experiment setting. Does the proposed method’s performance highly depend on hyperparameter settings?

review_title: Review

final_score: rating: 6: Marginally above acceptance threshold
review_content: Motivated by the issue of RL policy generalization, this paper explores improving generalization via a learned contrastive representation to embed states in, rather than through data augmentation or regularization alone. They demonstrate that their approach of policy similarity embeddings (PSEs) leads to improved generalization on several benchmarks, building on the method of policy bisimulation.

Strengths:
+ I think this approach is theoretically well-motivated, and the authors also give good intuition, especially in light of the latest work on self-supervised learning.
+ Experimentally, I thought the controls were well-chosen and illustrated the utility of their method on each task they considered (Figures 4 and Tables 1 and 3).

Weaknesses:
-	I would be curious to see how this approach works on more scaled up tasks with larger action spaces (Section 6.2 is in that direction, for instance), but the fact that it outperforms data augmentation techniques (including bisimulation transfer) in these simpler domains is promising.
-	In the jumping task, no intuition is given as to why PSE underperforms in the “narrow grid” (Table 1, first row) compared to bisimulation transfer, and it would be helpful to explain the failure cases of their method in more detail generally.

As it stands, I think the ideas of this paper are interesting and novel, but I would like to see tasks with larger action spaces or more naturalistic inputs than the ones considered here, and in the cases where the method does not perform well, documentation and potential intuition for why that might be. Therefore, I recommend a weak accept.
review_title: Theoretically well-motivated approach and promising empirical results

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