Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: This is an interesting paper that proposes abstractions based on return distribution similarities to be used as auxiliary tasks to aid in learning. The idea is quite promising and I think could open up new avenues for future research, but it does not appear to me to be ready for publication yet. In particular, although the authors claim to be distinguishing based on returns, many of the design decisions seem to implicitly assume determinism, or near-determinism (in particular, see point 3 and 10 below).
The theoretical results seem interesting, but I have some questions on their validity and clarity (see point 7 and 8 below).
The practical algorithm has a few issues that need clarification (see points 3, 4, 9, 11, and 12 below).

1. At the bottom of page 2, the authors write "we are the first to leverage return to construct a contrastive auxiliary task for speeding up the main RL task." This is not quite true, see [Zhang, A.; McAllister, R.; Calandra, R.; Gal, Y.; and Levine, S. 2020. Learning Invariant Representations for Reinforcement Learning without Reconstruction". arXiv preprint arXiv:2006.10742 .].
2. In section 2.2 you should cite [Taylor, J.; Precup, D.; and Panagaden, P. 2009. Bounding performance loss in approximate MDP homomorphisms. In Advances in Neural Information Processing Systems, 1649–1656.]
3. In line 6 of Algorithm 1, $y = \mathbb{I}[b(R_1) \ne b(R_2)]$ seems problematic for stochastic returns. In particular, if the number of bins is very large and the returns have wide variance, $y$ will almost always be zero.
4. In line 8 of Algorithm 1 is $\hat{w}$ also learned?
5. In the last sentence of the first paragraph of section 4.1, rather than comparing against regular bisimulation it seems more appropriate to compare to $\pi$-bisimulation from [Castro, P. S. 2020. Scalable methods for computing state similarity in deterministic Markov Decision Processes. In Proceedings of the AAAI Conference on Artificial Intelligence.].
6. In equation (1), the minimization appears over $f$, but $f$ doesn't appear on the RHS. Shouldn't the minimization be over $w$?
7. In Theorem 4.1 it's not clear what role $x_1$ plays in the result. Why is this third state necessary? It does not seem to show up in the proof in the appendix. Further, the proof in the appendix could do with some elaboration, as it's not completely clear how lemmas B.1 and B.2 result in the proof of Theorem 4.1. It would be better if the authors restate the theorem statement in the appendix and were more explicit about the connections. There are no page limits for the appendix.
8. How do we get a sense for how big $|\Phi_N|$ is? Couldn't it be as large as $N^{\mathcal{|X|}}$?
9. Typically auxiliary losses are combined with the main loss into a single loss, but in Algorithm 2 you seem to be updating them sequentially. Why? Does the order of update matter?
10. In Figure 1, on the left, the purple rectangle says "Value net". If so, are you really learning a distribution?
11. In Figure 2, why do some algorithms only have triangles and not learning curves? Also, it seems from that figure that "State SAC (skyline)" seems to outperform all others, including RCRL.
12. Is Figure 3 over a single seed?
13. In Definition A.1, bisimulation was _not_ introduced in (Jiang, 2018), it was introduced in [Givan, R., Dean, T., & Greig, M. (2003). Equivalence notions and model minimization in markov decision processes. Artificial Intelligence, 147, 163–223].


Minor comments:
1. It would be helpful if the authors specify how to pronounce RCRL. While reading the paper I was pronouncing it [like this](https://youtu.be/dQw4w9WgXcQ).
2. In the second-to-last sentence on page 1, should read "while ignor**ing** return-irrelevant features."
review_title: Official blind review

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

The authors present a contrastive auxiliary loss based upon state-action returns.

They introduce an abstraction over state-action pairs and divide the space of state-action returns into K bins over which the Z function is defined where Z(s,a) is distributed over K-dimensional vectors.  Given an encoding function, phi, and an input x, Z-irrelevance is defined as phi(x_1) = phi(x_2) when Z(x_1) = Z(x_2) which motivates the objective for Z-Learning: to classify state-action pairs with similar returns (within bounds) to be similar.  From this a contrastive loss can be defined (Return-based Contrastive RL, RCRL) where class labels are determined by Z-irrelevance sets encouraging state-action encodings to be similar when the returns are.  In the limit Z becomes the RL state-action value function Q.

The authors evaluate their approach on Atari (discrete actions) and the DeepMind Control Suite (continuous actions) across both model-free and model based RL algorithms against and in combination with other auxilliary losses including CURL (Srinivas et al. 2020). 

Strengths & Weaknesses:

Auxiliary losses have become an important component in RL for developing stable agents that can generalize well and form good representations.  In particular, contrastive losses have come into increasing use with growing literature around these methods and so I believe the domain area of this paper is relevant and of interest.  The authors do a good job of covering the recent developments of background literature in their related work section and grounding their approach with recent efforts undertaken in RL auxiliary losses, contrastive learning approaches and state abstraction/representation learning literature.

The approach is overall novel as many contrastive learning methods are defined against input data or downstream representations, whereas this work derives it's data from RL returns and creates a link between the representational landscape of the observations and actions and broad outcomes as they are valuable to an agent.  As the author's have framed the problem, I believe this approach is more powerful and also more tractable than something like reward prediction.  Intuitively the formulation seems solid to me since we often would like to understand not only when we're in a good state and taking a useful action but also, in general, what kind of properties state-action pairs with similar returns should have.  The authors do note that this may be learnable by temporal difference updates alone however, this approach aims to directly encourage the learning of this relationship and decouple it from the RL algorithm (where perhaps other things may be focused on such as planning etc.). 

One shortfall of this approach could be the available data itself as you'd rely on the policy to provide you with good samples for RCRL.  The authors indicate that they segment trajectories to ensure better quality positive and negative samples for learning however, it could be made clearer how much of a problem this can be.    This approach could possibly be combined with a self-supervised approach to alleviate these types of concerns.   It would also be nice to know the additional computational burden of RCRL and how this compares to other auxiliary losses.

The experiments on Atari & Control look solid and demonstrate that this method attained a stronger score both alone and when combined with CURL and good top performance on DeepMind control suite tasks when compared again to CURL and pixelSAC.  It might have been nice to see more comparisons or combinations with other contrastive methods that have had some success in  learning visual representations (SimCLR: Chen et al. 2020, BYOL: Grill et al. 2020).  The similarity analysis also provided some nice insight into the inductive bias induced by RCRL.

Overall, the paper is well written and has a clear layout.  The authors provide clear algorithms and figures and the content flows well from section to section.


Recommendation:

I believe that this is a promising and very active area of research and that this work makes the case for a solid new approach and a set of encouraging results to back it up. 
review_title: Review of "Return-Based Contrastive Representation Learning for Reinforcement Learning"

final_score: rating: 6: Marginally above acceptance threshold
review_content: ## Return-Based Contrastive Representation Learning for Reinforcement Learning
### Summary
The authors propose Return-based Contrastive Representation Learning (RCRL), a contrastive auxiliary learning task that guides the feature network to encode representation relevant to the task rewards. The experiment results show that RCRL helps improve two commonly used RL algorithm (RAINBOW and SAC) in low data regime. Additionally, RCRL can also be used in combination with other auxiliary tasks to boost performance.

Overall, the paper is well-written, the topic is relevant to the field and the approach is novel.

### Strength
- Theoretically-backed.
- The topic of representation learning is pretty relevant to the field now.
- The learned representation is "task-relevant", and therefore can achieve higher performance compared to other representation learning methods.

### Weakness
- The reliance on environment returns make the approach pretty susceptible to poorly or sparsely defined rewards. Specifically:
  1. The auxiliary loss does not work in sparse reward environments.
  2. In the task with dense but deceptive rewards, the representation may be biased toward representation that is not helpful in the long run.
- The improvement in continuous control tasks seem to be really marginal. Why is that?
- The learned representation may not be very general due to its reliance on return signals. Certainly, it can help achieve better performance when we only focus on a single task with a well-defined reward function. Yet, the representation may not be as useful when we considered some practical real-world settings that require policy adaptation and transfer learning.
review_title: Relevant topic with novel formulation

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

The authors propose the inclusion of an auxiliary task for training an RL model, where the auxiliary task objective is to learn an abstraction of the state-action space that clusters (s,a) pairs according to their expected return.  The authors first describe a basic abstraction learning framework (Z-learning) followed by the extension to Deep RL as an auxiliary task (RCRL). The authors present results in Atari (discrete action) building on Rainbow, showing an improvement compared to baselines on median HNS in the low-data regime, and results on DMControl (continuous action) building on SAC, showing similar or improved performance compared to baselines.

**Quality**

Overall I found the approach and results to be interesting and moderately compelling. At first glance the improvement is surprising, given that model-free Deep RL already needs to abstract the state space on the basis of returns even without an auxiliary task. The key appears to be the focus on sample efficiency in the low-data regime, where the task seems to improve non-local value signal propagation compared to a bootstrapped algorithm (particularly on Atari, note that in the 100k regime the model-free algorithms have not yet learned to play Pong). Since it is not clear that the algorithm will generalize to more data (it's easy to imagine that the abstraction task will hinder performance when the base algorithms become more finely tuned), I would like to see more clarification of the goal throughout the paper (e.g. "In the low-data regime, our algorithm outperforms strong baselines on complex tasks in the DeepMind Control suite and Atari games" in the Abstract), as well as a reference to the focus in the Conclusions.

In the low-data regime, it's also critical to justify this approach compared to a model-based alternative. On Atari the authors compare to SimPLe, but MuZero would be a stronger baseline.

Besides the empirical results, the authors also nicely provide a description of the Z_\pi abstraction and an error bound.

**Clarity**

I was confused by the description of the positive/negative sampling procedure in 4.3 paragraph 2. Are segments temporally consecutive within a trajectory? If so, is it primarily a heuristic that they will "contain state-action pairs with similar returns" (i.e. couldn't a reward achieved mid-segment make this statement incorrect)? As I understand it, segmenting avoids the problem of determining bins on the return distribution a-priori, however it also seems like it will limit the agent's ability to cluster non-local (s,a) pairs with the same returns. It might also mean that the agent is learning to cluster temporally adjacent states in the underlying state space rather than similar returns.

**Originality**

The paper builds on existing work in the abstraction literature and auxiliary tasks for deep RL. The primary novel component is using a return-based auxiliary task. 

The Z_\pi abstraction framework also appears to be novel, although its closely related to existing abstractions like Q_\pi abstraction.

**Significance**

The RCRL model itself improves on existing model-free approaches and can be easily incorporated into many model-free architectures, although it seems unlikely to beat a strong model-based baseline like MuZero in the low-data regime. 

The description of the Z_\pi abstraction, and the exploration of return-based auxiliary tasks in general, could prove more significant in the long term.

**Pros:**
- The model improves performance in the low-data regime over existing model-free baselines
- The model can be easily added to many existing architectures
- Description and theoretical results on a new type of abstraction

**Cons:**
- The paper needs some more clarity around the focus on low-data / sample efficiency and how applicable the model is to higher data regimes
- Unclear if the segment-based sampling strategy is clustering (s,a) pairs with similar returns or just states that are nearby in the underlying state space
- The model seems unlikely to improve on a stronger model-based baseline in the low-data regime

review_title: Nice results, but focus on the low data regime should be clarified

Meta-review: RCRL is return-based contrastive learning for reinforcement learning, where the label is whether two samples belong to the same return bin. The reviewers found this to be a well executed paper with good theoretical and experimental results.
Decision: Accept (Poster)