Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: summary: This paper proposes an algorithm learning a kind of universal value function (UVFA), so that given a buffer with a good coverage of transitions it returns the optimal behavior for any reward function or any goal state by computing a simple average.

The idea comes from the fact that the Q-function for a given policy $\pi$ -- that induces a reaching probability from any ($s,a$) $M^\pi(s, a, s', a')$ -- can be expressed as the scalar produce $Q^\pi(s,a) = \sum_{s',a'} M^\pi(s,a, s', a') r(s',a') = \langle M(s,a); r\rangle$.
Then, if we learn $F(s,a,z)$ and $B(z)$ such that $M^{\pi_z}(s,a) = F(s,a, z)^T B(z)$ we obtain $Q^\pi_z = \langle F(s,a, z)^T B(z); r\rangle = F(s,a, z_r)^T z_r$ where $z_r = \langle B(z); r\rangle$.

The greedy policy given such a Q-function should maximize (over action) $F(s,a, z_r)^T z_r$. Now, there is a fix point {$F^*, B^*, \pi^*$} such that for all $z$, $Q^{\pi^*}(s, a) = F^*(s,a, z)^T z$, which coincides with the optimal Q-function for the given $z$. 

The paper provides an algorithm that learns these optimal $F^*$ and $B^*$. Then, given any reward function, they compute $z_r = \langle B^*(z); r\rangle$ and return the optimal Q-function $Q^*(s,a) = F^*(s,a, z_r)^T z_r$.

main_review: The idea is brilliant, and I think will have a strong impact on the topic of UVFA. 

However, reading this paper was quite painful, in the sense that it is really badly constructed. A lot of quantities are introduced before the notation is explained, or ideas spoiled before the content makes it understandable. I agree that in the abstract and introduction the authors don't have the choice and must explain the idea before having the mathematical notation and tools. But --in my case-- I finally understood what the authors were talking about  at line 183 with the "sketch of proof" (page 5). 
I would suggest a lot of improvements regarding the construction of the paper, but to rather summarize:

1) separate introduction and related work, and move the related work later in the paper (for ex before or after experiments)
2) improve the introduction by a better explanation "with hands" of what $F$ and $B$ represent (future and past), how to use $B$ to compute $z_r$ and then $F$ to get the optimal policy. 
3) in problem and notation, give the problem after the notation.
4) right after the notation, briefly explain the idea like I did in the summary of this review.
5) put elements of your proof in the paper, so that you prepare the reader to understand the TH1
6) ends up with TH1

In the current state, we don't understand the quantities before TH1, then we don't understand TH1, then there is a long explanation that we finally understand by reading the short "sketch of proof". 

Another problem of this paper regards the experiment design: did the author used more than one seed? if yes how many? (especially for the pacman expe). The pacman results are promising (FB works better than HER), it would have been great to have more examples of such complex environments.

I think this paper is already nice and important to be communicated, but I would be a bit disappointed because of its actual form (difficult to read), so I am on the fence, marginally above the acceptance.

limitations_and_societal_impact: This paper has no societal impact.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This paper considered the representation learning in reward-free MDPs to encode the optimal policy given any reward function without planning. The proposed unsupervised learning algorithm FB offered two features $F$ and $B$ that encoded the information of the transition to extract the optimal policy. They showed that as long as $F^\top B$ approximates the probability of visiting one state-action to another, then we can extract the optimal policy given any reward function directly. Motived by the nice theoretical properties of $F$ and $B$, they showed that the unsupervised learning algorithm can achieve the same level of learning compared with DQN, DQN with HER, etc. on several benchmark environments. They also incorporated some engineering tricks to make the theory-boosted algorithm more effective.
main_review: Novelty: The paper proposed a novel representation learning idea to learn the forward-backward features $F$ and $B$.  The authors constructed $F$ and $B$ according to a novel observation that $F^\top B$ should approximate the probability of visiting one state-action to another. They also proposed new unsupervised learning method to train $F$ and $B$ based on a new density measure. 

Quality & Clarity: The claims are clear and well-supported to me. The proofs are clean and correct. The organization is okay, but some statements in several sections are obscure and took quite a while to understand. For example, the authors introduced their ideas in the related work section, in a fashion of story-telling. This type of writing may fail to highlight the critical points of the paper, and make the readers confusing because they are not comfortable with the notations. 

Significance: The paper studied a very important problem for both theoretical RL and empirical RL: how to “summarize” an environment, so as to remove the useless information and keep the important part. The FB representation provided valuable insights on this problem. The experiments on Maze and Ms. Pacman also verified that FB is comparable to the classic value-based methods, but FB is more flexible to random goals.

limitations_and_societal_impact: The limitations have been addressed in the paper, including the lack of exploration, the large variance, etc. The lack of exploration can be a serious problem, because rewards in real-world are mostly not random. Perhaps this problem is related to some techniques in batch RL. Moreover, it would be better to show more experimental results, such as in more realistic environments, or compared with more goal-oriented baseline algorithms.
review_title: null

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: summary: This paper proposes an algorithm that is able to learn the inherent structure of a reward-free MDP in a compact form, which can quickly give a nearly optimal policy after the reward function is specified. The algorithm is developed through a rigorous theoretical analysis and the experiment also shows comparable or even better performance than the algorithm trained with specified rewards.
main_review: Overall, I appreciated the novel and rigorous theoretical framework developed for learning MDP representation in this paper. Meanwhile, the experimental result looks promising and the extended analysis in appendix looks thorough. However, the writing of this paper may need some improvement since currently the ordering of some content makes it harder to understand, details in "Suggestion on writing" part.

Meanwhile, I also have the following concerns and questions.

Questions:
- On page 6, why it is okay to write $z_R$ as stated in Eq. (9) when the reward is known explicitly? Is Eq. (9) consistent with the definition of $z_R$ in Eq. (8)?
- In the experiment for FetchReach and Ms. Pacman, the behavior policy becomes adaptive, which implies a changing distribution $\rho$; meanwhile, the theory uses a fixed distribution $\rho$. Is there any intuition on why a changing distribution $\rho$ is better under this scenario?
- Does the distribution of $z$ play any role in the theory development? It seems $z$ is always fixed in theory while the distribution $\nu$ of $z$ is very carefully designed in experiment.

Suggestions on writing:
- At page 5, it may be better to explicitly point out that Theorem 6 in [BTO21] gives the update rule in Eq. (7).
- It may be better to move the explanations about $\mathscr{L}$ and $\mathscr{L}_{\mathsf{reg}}$ at Section D.3 to Section A right below the full Algorithm 1. Or more aggressively, it may be better to move all of these into the main text of Section 4, which can help a lot more in understanding the algorithm than reading the paragraph description in current Section 4. This extra space can be saved from moving some figures and experimental results into the Appendix.
- In Remark 4, page 17, it may be better to say explicitly that $B'$ being $L^2(\rho)$-orthonormal means $\mathbb{E}_{(s, a)\sim\rho}\left[B'(s, a)B'(s, a)^\top\right]=\mathbf{I}$ if I understand correctly.

```
[BTO21] Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal-dependent values: A mathematical viewpoint. arXiv preprint arXiv:2101.07123, 2021.
```

---

Good paper. My score remains the same after rebuttal.
limitations_and_societal_impact: The authors adequately addressed the limitations of their work in Section 4 by pointing out the potential problems of their algorithm. The potential negative social impact is not addressed, which will generally be all potential malicious usage of any reinforcement learning algorithms.
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: This paper proposes a new representation in MDPs such that it can be learned using transition samples and used to compute the optimal policy. I am not convinced by the motivation of such a new representation since no computational or statistical efficiency results are established and thus suggest not accepting the paper.
main_review: This paper proposes a new representation in MDPs such that it can be learned using transition samples and used to compute the optimal policy. 

My main concern is that such a representation might be too restrictive for learning unknown MDPs. As one can always learn the transition kernel well without knowing any reward signals and later compute an optimal policy by a dynamic program procedure using the learned model and reward function provided a posterior, the proposed new algorithm just provides an alternative which might save some computation in the second step while not sure adds how much restriction on the first step. For the tabular MDPs, the dimension d of is already |SA| as shown in the appendix and I am not sure how large the computation cost can be to learn the map F: S $\times$ A $\times$ Z $\rightarrow$ Z and B. To me, it feels equal to or more expensive than learning the transition kernel matrix. Also, it is not clear whether this new representation can help improve the exploration efficiency in the reward-free exploration phase, which is usually the result of many other reward-free works, e.g., see [1]. 

In all, I am not convinced by the motivation of such a new representation since no computational or statistical efficiency results are established and thus suggest not accepting the paper.

[1] Jin, Chi, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu. "Reward-free exploration for reinforcement learning." In International Conference on Machine Learning, pp. 4870-4879. PMLR, 2020.
limitations_and_societal_impact: yes
review_title: null

Meta-review: The majority of the reviews agrees that the paper presents an interesting new approach and votes for a clear accept. Accordingly, I recommend acceptance. As most reviewers also pointed out a few issues with the presentation and organisation of the paper, please take into account these comments when preparing the final version. Finally, one reviewer was not that completely convinced by the suggested representation. In the discussion (s)he said (I think expressing the point better than in the original review):

"My concern is that the representation carries too much information. To be specific, (please correct me if my statements are wrong), in their tabular MDP example, the function $F$ has a large continuous input space (dimension>SA) and thus requires function approximations and large computation power. In contrast, learning the MDP only involves $S^2A$ exact parameters, and the planning under a given MDP and reward function is usually not the bottleneck (using dynamic programming and given a discount factor in (0,1)). This issue can be even worse for general MDPs, and the paper seems to not discuss too much on the approximation power of such a representative for other MDPs."

Although I understand that the paper suggests not to learn a particular MDP but a range of tasks, the question whether the arising complexity of the final problem is reasonable when compared to the complexity of a set of single tasks is I think a valid one and should be discussed.
Decision: Accept (Poster)