Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper extends the idea of set transformer to motor control, so that an agent can perform well even when its sensory input channels are shuffled and partly presented.
The performance is verified in cart-pole, ant, pong, and car racing tasks by using evolutionary optimization and behavioral closing of pre-trained policies.

main_review: Originality:
The present result is an extension of the set transformer (ICML 2019). Although the title includes "reinforcement learning", actual learning takes evolutionary optimization and supervised learning (behavioral cloning).

Quality:
I wish the role of the previous action input to the attention network is better discussed. My guess is that the network compares the changes in the observations and action taken to infers which channel carries the critical information for control. It would be interesting to analyze which channels actually gathered attention (e.g., the patches including the car and road boundary?) and how that is affected if the past action is not presented.

Clarity:
The math is clearly presented.
In equation (1), most RL folks imagine Q as action values and V as state value, so their meanings as Query and Value may be better explained earlier.

Significance:
Although the architecture is reported to be inspired from sensory neurons, it has only remote link to neuroscience.
Although not tested in this paper, permutation-invariant feature would be helpful in the context of multi-agent learning, where an agent need to work with a variable number of agents coming and going in random order.



limitations_and_societal_impact: These are well considered and documented.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This work investigates the problem of shuffled or dropped out sensory inputs for reinforcement learning agents. The authors propose AttentionNeuron, a permutation-invariant attention layer operating on neurons which independently process the input currently present at its location. Each neuron utilizes its past context and the past action to interpret the seen input, and the attention mechanism combines the outputs of each neuron to generate a globally-aware embedding to be used for a standard policy architecture. The authors show that (1) AttentionNeuron is robust to shuffled and noisy observations on CartPole, (2) performing behavior cloning with an AttentionNeuron architecture can yield a policy robust to shuffling on Ant, (3) AttentionNeuron is robust to patch-level shuffling and dropouts on Pong, (4) AttentionNeuron improves robustness to drastic background changes on CarRacing, and (5) AttentionNeuron is robust to shuffling within a single episode.
main_review: Pros
1. Studies novel and biologically inspired permutation invariance question in the context of RL
2. Experiments evaluate on several diverse environments (CartPole, Ant, Pong, CarRacing)
3. Writing is generally clear

Cons
1. It is not clear to me that permutation invariance is a practically significant or useful property
2. Minimal ablations/comparisons to other methods in the experiments section

At a high level, I think this work is promising: it investigated a little-studied problem, permutation invariance to sensory observations in RL, motivated by observations in neuroscience regarding malleability in human brains. While it is not clear to me why this is a useful property from an engineering perpsective (i.e. as to why this can't be mostly tackled by the sensor interface), I think the direction studying biologically plausible RL/models is exciting. Furthermore, I think additional practical benefits gained as a result of permutation invariance (such as better generalization) are particularly exciting/important.

My main concerns are about the lack of comparisons/ablations to other methods, as the authors propose a method with several components but only compare to a feedforward network. I don't necessarily believe that my choices below would outperform the proposed method, but I think they would strengthen the paper and better contextualize the presented numbers. To propose a non-exhaustive but simple set of choices:
1. Transformer/vision transformer: a transformer operating directly on the sensory inputs/image patches without positional embeddings is permutation invariant in the same sense as AttentionNeuron, but is a simpler approach (i.e. AttentionNeuron without the LSTM input cells). I think this ablation would make clear whether or not the recurrent layers are necessary, and if it is important that the bottom-level cells "orient" themselves,or if simply the upper-level attention layers are capable of making sense of the disordered information. I also think a comparison with positional embeddings would be useful, as vision transformers currently represent a state-of-the-art in image-based processing.
2. LSTM after image processing: similarly to my above comment, this would clarify if this recurrent reorganization should be done at a low-level or a high-level. I think this is also biologically interesting, and practically it would be less computationally expensive since there wouldn't need to be an LSTM at every input location
3. Other generalization methods: as I mentioned above, I think a particularly exciting experiment is Section 4.4/Table 4 because it promises that generalization could be improved via permutation invariance. However, the experiment only compares to a minimal baseline. I think comparing to other representation learning methods, such as bisimulation [1] or contrastive learning [2] (or others), would make this experiment more exciting.

Overall, I think the direction toward permutation-invariant representations for sensory observations is underexplored and novel in the RL setting. Also, I think the authors investigate a few different interesting properties of the method, though I think the baselines are weak. As a result of the aforementioned weaknesses, I think this work currently does not meet the bar for publication. I am open to changing my score if the above issues are addressed, particularly if the comparisons to the transformer and high-level LSTM are included.

=====

References

[1] Zhang, Amy et al 2020. "Learning Invariant Representations for Reinforcement Learning without Reconstruction."

[2] Srinivas, Aravind et al 2020. "CURL: Contrastive Unsupervised Representations for Reinforcement Learning."

=====

Updates

I am satisfied with the authors' responses and additional baselines, therefore I am increasing my score from 4 to 6.
limitations_and_societal_impact: Yes, the authors adequately addressed the limitations and potential negative societal impact
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This paper presents a sensory neuron architecture as a layer in a neural network trained for reinforcement learning and claims that this sensory neuron makes a network generally permutation-invariant. There is an extensive related-work section covering examples of sensory modality adaptation in neuroscience, cellular automata, meta-learning, and attention, setting up the presented approach which uses collective intelligence of only distributed modules with only local awareness, self-organization through attention, and adaptation of an existing policy to be permutation invariant. The paper then presents a mathematical formulation for the AttentionNeuron architecture, including its temporal memory, to be used as the first layer in an RL agent. The paper next tests its approach on four different tasks, some with multiple environments and concludes with some discussion of the results which appear to be robust to permutation of the input vector unlike the baseline feedforward network. 
main_review: *Originality:* I think this paper’s originality comes from the fact that it is a novel combination and then application of well-known techniques and concepts. The paper notes that Transformer’s attention mechanism is known to be permutation invariant and applies it here to reinforcement learning, which is the primary novelty I see in this work. The paper suggests an angle of collective intelligence due to the AttentionNeuron architecture, where each of N sensory neurons sees one of each of N observation components. It seems that the transformation from input to sensory neuron layer is a generalized permutation matrix, and the paper’s point is that this can be any generalized permutation matrix as long as the magnitude of each column remains the same (I think; I may not have all the details totally right). The connection to collective intelligence seems limited to a high-level abstraction, but I do think it’s an elegant and unique approach to permutation invariance whether or not it would be considered collective. 

*Quality:* I think the methods and results of this paper are quite well-conceived and -executed. Though the method seems complex at first and to some degree needlessly complex, the discussion of design choices such as temporal memory and reduced network size is helpful. The solution is actually quite thoughtful, and the issues are more of presentation than content (discussed below). The experiments are impressive - the approach is used successfully in four different testbeds, some with multiple environments. Though there are some caveats with each one, particularly Atari Pong, the core message of permutation invariance seems to persist through the results. In my opinion, the paper is really strengthened by forgoing complicated cuts of data and various statistics to just show performance on a large range of testbeds. The only improvement may be ablations or analysis to rule out confounds.

*Clarity:* This paper is somewhat difficult to comprehend. Section-specific critique: 

The related work is quite ambitious. I think all but the attention section could be collapsed into a section about permutation invariance and *maybe* collective intelligence, though as I said in *Originality*, that connection seems vague. Many of these are only related loosely, and the authors say as much at the end of each paragraph. This paper presents an interesting layer architecture that achieves a clear and concrete goal, and related work about that goal and similar architectures is enough in my opinion. This related work section is confusing and to some degree feels overpromising despite a paper that doesn’t need to overpromise. This makes the paper difficult to understand. 

In the methods section, the notation is difficult to follow, partially because it is mostly in paragraph form. Figure 2 is quite helpful, though the 1:1 nature of the observation components to sensory neurons could be visually clearer, but the rest of the text on page 4 is dense with words and notation combined. What makes it more frustrating is that it is notation, not derivation or theory, so it shouldn’t be difficult to follow; a diagram or creative visual aids for the central definitions might help. 

Finally, the results are very hefty. That’s one of the biggest value-adds of this paper, but it’s also overwhelming because it’s presented as one long laundry list. I realize that each testbed is very different so it may be hard to have a claims-driven structure to this section, but it would be nice to see unifying stats or even unifying qualitative analysis earlier than Section 5. 

*Significance:* I think the ambitious and extensive related work section causes this paper to feel like it’s overstating its significance; I don’t think neuroscience is really a related field at all, nor am I convinced that this is an example of collective intelligence, nor do I know why it would make a difference. It’s also hard to tell if permutation invariance itself is significant - I think the significance of the main contribution comes from the fact that it’s 1) an elegant architectural approach (which builds on Transformer), unlike existing permutation-invariance works that require specific input formats and 2) application to reinforcement learning. If the language focused on what this paper is really doing, as opposed to big ideas that sound kind of like what it is doing, it would come across better. Nonetheless, I do think that this result is significant. The applications listed in section 5 are a good taste of what this method can do, and I suspect there are many more, and more immediate, applications that could be explored from this work.
limitations_and_societal_impact: The limitations receive some discussion, though I think the paper would benefit from more beyond resource usage. I also don’t think the lack of invariance to action permutation is as significant as a lack of invariance to observation permutation would be, though this could be a limit of my own knowledge. Ablation experiments may have been helpful to yield this section. The societal impact section is considerable and far-reaching, and I think it briefly hits on many important points. I want to clarify that I think the way this paper is written calls for exactly this type of societal impact section, so the authors have done a good job. That said, I personally prefer depth and engagement with societal implications rather than a broad laundry list, and if the paper had a narrower focus, Section 6 could focus on just the implications of permutation invariance (i.e. some of its main applications) and/or RL agents.
review_title: null

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: summary: This work unites previous research on Transformers and concepts from cellular automata to offer a novel model for learning global policies from local arbitrarily ordered variable-length sensory inputs. The authors evaluate their model using four standard RL tasks and show that their model, compared to the state-of-art, is robust to perturbations and occlusions of the inputs and also is significantly more effective for out-of-distribution generalization
main_review: This work provides an interesting and highly valuable method for building RL models robust to various types of noise and capable of out-of-distribution generalization, which is a highly important and widely pursued goal in RL. To attain this goal, the authors propose a novel framework based on Set Transformers (a Transformer-based method for permutation-invariant neural-network analysis of datasets) and also based on concepts of cellular automata. The authors rigorously examine the properties of their new framework using a diverse set of standard RL tasks. As a result, the authors show impressive capabilities of their model to perform robustly in the presence of large-scale occlusions, noise, and out-of-distribution versions of the tasks. With that said, the text could use numerous clarifications and a comprehensive methods section, which would make the paper accessible for the broader NeurIPS community. Overall, the novel model, sound methodology, thorough testing, and highly significant results presented in this paper make it a valuable contribution for NeurIPS.

Specific comments:

Although the paper seems to have all necessary details to be understood, these details appear to be scattered throughout the text (and throughout the referenced literature) which may require reading the paper 2-3 times to cover all the important details. To this end, I suggest making use of unlimited supplementary space to write comprehensive methods, which – alone – would be sufficient to reproduce the model and the experiments.

The paper may also benefit from adding the details of the framework to the main text. Notably, the Set Transformers used here are much less known than the original Transformer framework, so it may be unfair to expect the reader to know its details or to be motivated to read a separate paper on it. As a good example, the Set Transformers paper itself starts with a brief yet sufficient introduction of the Transformer framework and explains the design choices for the Set Transformer framework immediately upon their appearance. One of such design choices – not explained in the current paper – is that SensoryNeuron, just like Set Transformers, attends the input data (Keys, Values) to a (smaller) set of separately learned/handcrafted Query embeddings – unlike the usual Transformers attending the elements of input data to each other. It is important to explain that learnable/predefined Query embeddings do not depend on the input and decrease the algorithm’s complexity from quadratic to linear.

The mathematical notation here can also be simplified, as it appears to have multiple variables which are not used. For example, the matrices Wq, Wk, and Wv were originally introduced in Transformer to be used in multi-head attention, where the different versions of these matrices were applied to the same Queries, Keys, and Values. It appears that the framework presented in this submission does not use multi-head attention, so these matrices can be sacrificed for the clarity of how the main model is explained. The matrix Wk, in particular, can be described together with the function Fk as both of them are implemented using neural networks – stacked with each other and thus forming a single neural network. The same goes for Fv and Wv – both being implemented as the identity matrices and thus unnecessary. Decomposition of Q and Wq is justified as it decreases the overall number of parameters and because Q is not learnable – but this distinction does not need to be immediately introduced. Focusing rather on the roles of Q, K, and V would make the paper more accessible for the broad audience.

Finally, I found the neuroscience-oriented introduction to this paper questionable. While the paper presents interesting, significant, and valuable results regarding the robustness of the algorithm to various types of noise, taking advantage of additional data, and out-of-distribution generalization – these results are summarized only at the end of the paper. At the same time, a reader who starts reading the text from the beginning may falsely think that the purpose of the paper is to learn from/replicate a biological phenomenon. Moreover, the introduction – in light of the details of the proposed model – may leave a reader with the impression that the brain information processing is permutation invariant, like in zero-shot learning, whereas it’s not. The neuroscience papers cited in the introduction, such as reverse biking or learning to “see” using an electrode matrix, clearly mention the significant time tolls of acquiring such skills. Although the presented work may have future neuroscience implications, the line of the argument currently presented here does not make a solid connection to neuroscience and, as such, needs to be revised/shrunk. Instead, the paper can start from the useful properties of the model, as outlined in the conclusions. The same goes for the abstract. The reclaimed space can be used for a detailed description of the SensoryNeuron, which is a central part of this work.

I would like to highlight here that all my suggestions above do not devalue the significance and the validity of the results. They are intended to highlight that a thorough explanation of the design choices, a broader overview of the Transformer, and starting from the robustness/generalization properties will make the paper a more accessible read for the broad NeurIPS audience. Even in the current format, this work is publishable in NeurIPS.

Minor comments:

-Why is the pass-through function Fv needed at all? Could you please briefly comment on the intuition why Fv, normally being a learnable parameter in Transformers, does not require learning here?

-Could you please also compute the statistical significance for the RL results? Given the huge standard deviation in the CartPole task, the SensoryNeuron’s apparent slight deficiency may not be statistically significant, which will strengthen your claim.

-Can you please explain why you spatially arrange the messages Mt in the AtariPong task? Naively, it would seem that, as you permute your sensory neurons, the 2D arrangement of their messages will be also random and the convolution on top of these messages will learn the same as the feedforward connectivity?
limitations_and_societal_impact: The authors have adequately addressed the limitations and potential negative societal impact of their work.
review_title: null

Meta-review: I thank the authors for their submission and active participation in the discussions. Reviewer dfwx has pledged that with the revisions by the authors, they would raise their score. Since reviewer dfwx has not done that yet and did not actively participate in the discussion with the other reviewers, I assume a higher rating than the one currently recorded. Taking this into account, it seems reviewers unanimously agree that this paper is worthy of publication. In particular, reviewers appreciated the diverse set of environments used for evaluation [rWCR,TUZP], that is well executed [Bz2K] and that it presents an original synergy between existing techniques [Bz2K,TUZP]. During rebuttal and discussions, reviewer rWCR's concerns regarding baselines were addressed and reviewers TUZP and Bz2K agreed that the paper is improved based on the author rebuttal. Thus, I recommend acceptance and encourage the authors to further improve the clarity of their paper based on the reviewer feedback.
Decision: Accept (Spotlight)