Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper proposes a representation learning method that leverages unsupervised signals like flow and forward models and the constraints between them and apply it to the state representation problem of RL. 

The method is adding a number of auxiliary rewards to the torso which flow and a latent transition model that are constrained to agree with a contrastive loss. The flow has an architecture inspired by FlowNet and trained with warping reconstruction and a spatial smoothness term. A latent transition model is trained contrastively. Then constraints are added to make the latent of the next step agree with both of these predictions. 

The method is run on 13 Atari tasks and compared there with competitive baselines and it largely outperforms them. An ablation without baselines is performed a few dm control suite tasks  
main_review: strengths:
* the method does outperform competitive baselines on challenging atari games
* the paper is well written and the method seems good.
* results are shown with a fixed set of loss weights which shows robustness.

weaknesses:
* the method is very complex and hard to reproduce. Even if the code is made available, building on top of this method will not be trivial.
* the ablation experiments are done on domains that favour good flow representations and lack baselines (dm control)
* the crux of the state-representation problem is that it should make credit assignment easier but the hard exploration games are either missing from the experiments or improvements seem marginal.

remarks:
* i think structured representation is a misnomer here. This is because the image lattice is the weakest form of structure you could have and because structured representations in the literature refer to objects and parts. So I recommend changing the title to reflect that.
* i can't seem to find the L_rl loss in your text. Is it SAC ? I guess it's just above the loss functions where you hint that both DQN and SAC where used somehow
* The method seems to struggle with pong ? Is this because the RL loss in this setting and these hyper parameters is washed away by the auxiliary loss gradients which have a hard time converging ? 

summary_of_the_review: The method seems reasonable but the approach is very complicated and there are some lacking experiments that would make the claims more compelling. If the community is to be convinced to embrace such a complicated approach experiments need to be more compelling i should think.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: In this paper, the authors focus on the problem of representation learning for deep reinforcement learning. To this end, they propose an approach to learning structured representations via establishing flows between latent volumes. Similar to SPR, they predict future representations with a transition model conditioned on actions. They evaluate their approach in two domains: Atari and DeepMind Control Suite and show improved performance over the state of the art. 
main_review: Strengths:

1. Quantitative results are promising. The proposed approach clearly shows promise over prior approaches.

2. Evaluation is thorough. The results hold across different environments and even different agent setups (Rainbow and SAC)

3. The paper is well-motivated on a conceptual and mathematical level.

Weaknesses:

1. The technical novelty is limited. The approach is very similar to SPR with the addition of latent variable warping.

2. It's difficult to conceive how this method could apply to different input modalities (i.e. audio, text, joint positions). The flow-based approach is very much tied to image inputs. 

3. While there is improvement in quantitative performance over baselines, the margin of improvement is fairly small.
summary_of_the_review: The authors propose an approach to self-supervised learning for reinforcement learning which is very similar to SPR. The key contribution is the addition of a flow-based approach to predicting future latent vectors. Because of this, the generality of the approach is limited to image inputs. However, the results are promising and hold across different domains and agent setups.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper focuses on self-supervised visual representation learning for RL applications. As opposed to previous literature that focused on learning global representations for the current observation, the authors propose learning "structured" representations based on flow maps that encode local structure. To do so, they combine:
- a self-supervised flow model (from previous literature), trained using an image reconstruction loss;
- a classic action-conditioned 1-step transition model using cosine similarity in latent space
- a flow based warping (also from previous literature) to provide a second next-step latent prediction task.
They evaluate their method on various standard RL tasks in two domains (Atari, DM Control Suite) and show improvements over SOTA in ~half of the tasks.
main_review: The paper is very well-written and referenced, and the figures are quite clear (especially Fig 2 which gives a great overview of the method). I find the idea of leveraging self-supervised flow models quite interesting as an alternative to brute-force next-step prediction.

My main concern, however, is that the results the authors chose to put forward are not as good as I'd have liked, which makes me doubt the overall significance of the method:
- in Atari, the proposed S³R performs significantly better than SPR in 12/26 tasks, and underperforms significantly in 7/26 tasks;
- in Control suite at 100K steps, S³R significantly outperforms DRQ in just one task out of 6, and underperforms significantly in 2, and both are more-or-less tied at 500K steps;
Since S³R seems significantly more complex in terms of implementation and computational cost, it seems unlikely that a practitioner in the field would bother compared implementing it rather than using a much simpler variant.

Some remarks/questions:
- I'm not aware of a definition of "structured representation", but I'm not sure I'm convinced by the claim that S³R's repr is more structured  because its features translate to flow maps. I also don't think it's necessary for the paper to make this claim.
- Introduction: "Learning visual features from raw pixels only using a reward function leads to limited performance and low sample efficiency."  I would tame down/remove that claim. I don't think there's any consensus on that, and work such as  https://arxiv.org/pdf/2104.06294.pdf are able to improve performance and sample efficiency without relying on an explicit visual features representation learning scheme
- I was surprised by the results of Table 3 (that F+W > P, and that F+W+P > F+W), as I would have expected next-step latent prediction to be a somewhat easy task. Do the authors have any further insight on this, eg on how (dis)similar Z^w and Z^p are, maybe in terms of an auxiliary stop-gradiented Sim operator?
- Following up on that, could using both Z^w and Z^q as input to the RL head bring an improvement?
- Notations: the paper (eg fig 2 and the paragraph at the top of p5) seems to imply that a_k is computed from I_k to I_{k+M}, should this be I_{k-M} to I_k instead?
summary_of_the_review: The paper combines existing self-supervised architectures (self-supervised flow estimation and InfoNCE-like similarity loss) to learn visual features in a rather novel and intriguing combination, but the overall increment in RL performance does not seem to justify the added complexity compared to much simpler counterparts.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The paper proposes a new self-supervised representational learning for reinforcement learning tasks, and demonstrated superior performance in the sample-efficient setting against states-of-the-arts models. The major idea is to force neural networks to encode optical flow between consecutive frames and impose regularization that the local embedding predicted by warping and a state transition model (given action) are both similar to the extracted embedding from the next frame.
main_review: The motivation is justified because indeed previous literature has shown that when humans play Atari games, representation of objects are important for efficient learning. Learning global representation may lose detailed information of objects. Although the current work does not attempt to explicitly model objects, the flow information likely captures some object information implicitly.
I find the results generally convincing and the idea appears novel and interesting.
However, I am a bit worried why the internal flow loss in equation 4 will help the network extract useful information. If my understanding is correct, the flow for warping internal representation is the same flow as the one for warping image (equation 2). How can we guarantee this the extracted query feature does not simply become a fixed linear transformation of the local image patch? As long as the representation in equation 1 can be learned to predict external flow that satisfies the image warping cost, it seems to be able to guarantee the cost in internal flow to be low as well, even if query feature is a simple linear feature of local image patch.
Although the paper has demonstrated good performance in the RL task, it may be interesting, if possible, to provide some visualization of what kind of local feature is actually learned, in the appendix.
summary_of_the_review: I think the paper provides an interesting new objective for encouraging the representation by self-supervised learning also captures important local information for efficient learning of RL tasks.
review_title: null

Meta-review: After carefully reading the reviews and the rebuttal, unfortunately I feel this work is not yet ready for acceptance. 

I want to acknowledge the effort put in the rebuttal for this work and I think all the changes greatly increased the value of the work. However, I feel that the work could greatly benefit from running on a different domain where the gain is more considerable. My worry is that the complexity of the method compared to the relatively small improvement (at least as perceived from the current results) will reduce considerably the attention the work will receive from the community (unfairly so). 
Or some of the analysis and ablation done (e.g. the flow visualization) which are now in the appendix could be brought in the main manuscript to be able to drive the message home. An understanding of the impact on the accuracy of the flow model on the overall performance (which as pointed out by reviewer aHc1 is a really hard task in a more natural context). Or maybe a 3D visually complex environment is exactly where this method will shine as flows are more complex and hence more informative. 

Overall I think this is solid work, but I feel it does not manage to convince the reader of the significance of the proposed approach. And hence, if published in this form, I feel it will do a disservice to the work, as it will not receive the attention it merits from the community.
Decision: Reject