Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper proposes the Composable Perceptual Schemas recurrent architecture. This recurrent architecture is decomposed into $n$ LSTMs, each of which the authors call a _schema module_. At each time-step:
1.  a recurrent encoder produces a set of $m$ feature vectors of dimension $d_z$, one vector for a different spatial location
2.  the task information $\tau_t$, the previous reward $r_{t-1}$, and previous action $a_{t-1}$ gets concatenate with the previous hidden state $h_{t-1}^{(i)}$ of module $i$. The authors call this concatenation the _context_ $c_t^{(i)}$.
3. The update rule of each schema module produces the new hidden state $h_{t}^{(i)}$ from three inputs:
    1. The output of $f_{att}$: A single soft mask of length $d_z$ is computed from $c_t^{(i)}$ and is applied to each of the $m$ feature vectors.
    2. $c_t^{(i)}$
    3. The output of a multi-head attention layer, where the queries are the $c_t^{(i)}$, and the keys and values are the previous hidden states of the schemas, along with the null vector: $[h_{t-1}^{(1)}, ... h_{t-1}^{(n)}, 0]$

The main difference between CPS and the baseline recurrent architectures (AAA and RIMs) is in step 1 of the update, how $f_{att}$ is computed: instead of attending over the different spatial locations CPS attends over the same dimensions of all the spatial features.

The authors apply their architecture to three generalization scenarios:
1. " Can CPS enable generalization of memory-retention to novel spatial and temporal (spatiotemporal) compositions of object-dynamics?" --> tested on the Ballet task from https://arxiv.org/pdf/2105.14039.pdf. CPS performs slightly better than the LSTM when generalizing to different spatial compositions, and much better than all baselines when generalizing to different spatial and temporal compositions.
2. " Can CPS generalize sequential active perception of 3D objects to larger environments?" --> tested on the "Place X on Y" task from https://arxiv.org/pdf/1910.00571.pdf. CPS learns faster than the baselines.
3. "Can CPS generalize goal-oriented behavior and memory-retention to environments composed of longer sequences of observed object-configurations?" --> tested on a custom KeyBox environment written in the BabyAI framework. CPS performs similarly to the AAA baseline.





main_review: It is nice that this paper is motivated by representing schemas from cognitive science, and to my knowledge this paper provides a novel architecture attempting to do so. From what I understand, the architecture aims for each module to specialize to represent a particular spatial configuration of a scene, and different modules might represent different spatial configurations.

I have the following concerns:

1. It was difficult for me to take away from this paper a set of general principles that can be applied to future research. I would expect that the paper either (1) goes the scientific route: takes the best theory of schemas from cognitive science and show how their architecture is entailed by that theory, in which case the experiments would help us refute alternative theories of schemas proposed in cognitive science and give us insight about the nature of schemas from the abstract perspective, or (2) goes the engineering route: proposes an architecture that gives much better performance on the evaluation tasks, in which case the experiments would help us understand what engineering principles are useful to incorporate into future methods. This paper does not do (1), so I would expect it to do (2). However, the experimental results suggest that CPS is often similar to the LSTM baseline (section 4.1, 4.2), and gets the same performance as the AAA baseline (section 4.3). Furthermore, for the Ballet task, the paper does not compare against the method (HCAM) in the paper that proposed the Ballet task (https://arxiv.org/pdf/2105.14039.pdf), so it is unclear whether CPS offers insight on the best way to solve such a task.
2. It was not sure what to conclude based on the analysis in section 4.3.1 and Appendix C. For example, could it be possible that the plots of Figure 6b, Figure 8, and Figure 9 could have just as well have been generated by just looking at the random initial weights of the schemas? Given that the authors are motivated by modeling schemas that model independent aspects of the task, what measurement could the authors make to evaluate whether the schemas have specialized to learning something meaningful about different aspects of the task, in a way that would be distinguishable from the measurement take on randomly initialized schema modules?
3. I am unsure what to conclude from Figure 3c. It seems that having the ablations of having 4 subschema modules without attention and having only 1 subschema module with attention performs similarly to the full CPS. If this is the case, how would I evaluate the CPS as a contribution, if ablations that remove the core parts of the method (>1 schema, or feature attention) still perform similarly well? If the authors hypothesize that "We hypothesize that learning n subschemas, each with 1/n parameters of a single monolithic module, reduces overfitting and encourages the discovery of regularly occurring structures," does Figure 3c actually refute their hypothesis?
4. In various occasions the paper makes claims about the abilities of the CPS module that are not sufficiently tested. For example, the authors state in 4.1 that " Learning this task tests a recurrent architecture’s ability to recognize and maintain separate,
independent dynamics in an agent’s state representation" --> can the authors please conduct the experiment that tests whether the schema modules actually learn independent dynamics with an appropriate metric of independence? As another example, in section 3: " Since the feature-coefficients for the next time-step are produced with observation features from the current time-step, subschemas can dynamically shift their attention when task-relevant events occur" --> can the authors please conduct the experiment that evaluates whether this shift in attention is actually true? If there is a way to visualize how each schema module's attention to different features changes over time, and match that attention activation to semantically meaningful events in the episode, that would give much needed insight on how the method works. Another example is, in section 3: "To encourage specialization, each subschema has $\approx \frac{1}{n}$ the parameters a single module would have." --> can the authors please conduct an experiment that measures the degree of specialization with an appropriate evaluation metric, for various different numbers of parameters of the subschema? It is unclear to me that simply restricting the parameter count is enough to induce specialization.
summary_of_the_review: Overall, while this paper is tackling an important question, its execution could be substantially improved. It was not clear what the contribution of the paper is from a scientific or engineering perspective. The analysis could be more informative about what the modules are actually learning. The central claim of the necessity of feature attention and multiple schema modules seems to be refuted from Figure 3c. The paper needs to either temper its claims or provide empirical evaluation for its claims. More details about these concerns are in the Main Review. It could be that I may have misunderstood something, in which case I would happily engage in discussion with the authors. 

== AFTER REBUTTAL ==
I thank the authors for their response and additional experiments. I have raised my score to a 5, and I especially appreciated the clarifications and the qualitative analysis. After reading the rebuttal and the revised paper, one concern I still have is that it is difficult to glean a single coherent empirical message from the paper: on some tasks FARM performs the baselines (Figure 3, 4), on other tasks FARM performs similarly (e.g. Figure 5b, 7, 8). Perhaps this could be something that could be fixed through the writing, but overall the case for the following questions could be made stronger: (1) Why these three types of generalization rather than some other types of generalization? What is the underlying property that all these three types of generalization have in common, and why is this property important for AI research? (2) What is the explanation for why the proposed method must be the method for tackling tasks with this underlying property?

Regarding (1), it would strengthen the paper's overall message to identify a single unifying characteristic of the problems considered, because at the moment the paper is written in a way that highlights the disparate nature of these generalization tasks, but without an underlying theme connecting these generalization tasks, they seem to be chosen in an ad hoc way. Regarding (2), this could be addressed with stronger empirical results.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes Composable Perceptual Schemas (CPS), a modular state representation learning architecture for reinforcement learning that combines modular RNNs as in RIMs (Goyal et al., 2020) with a dynamic feature attention mechanism. It is hypothesized that this lets CPS exhibit multiple different kinds of generalization encountered previously in the literature: (1) generalization to environments given by novel compositions of known object motions, (2) generalizing active perception of 3D objects to larger environments, and (3) generalizing goal-directed behavior and memory retention to larger environments.

The idea of CPS is as follows:

* A recurrent observation encoder produces observation features for different parses of the input observation (here spatial locations in the image) to yield a feature matrix. Here each row corresponds to a different parse.

* A number of recurrent modules (here called "subschemas"), each process the observation features + context (action, reward, task encoding) and the previous state of all other modules using dynamic attention. From this, an updated state is computed using an LSTM.

* The dynamic attention for observations + context proceeds by scaling the observation features at all positions with a vector of learned attention coefficient. In this way, feature selection can take place (but shared across positions).

* The dynamic attention for other modules is similar to the implementation in RIMs using transformer-style attention.

CPS is evaluated on three environments and compared to several baselines. In general, it can be seen to exhibit the desired generalization behavior and outperform the baselines. Finally, some evidence is presented that subschemas are selective for semantically meaningful aspects of the input (eg. goal information, distractor information, etc.).
main_review: The focus of this paper is on addressing three particular types of generalization in RL as listed above. A commonality between them is that test tasks are novel compositions of regularly occurring "structures" the agent experiences during training. In the first case, these are the object-motions, and in the second and third case the visual objects. The difference in generalization behavior in the second and third cases is less clear to me since both can be viewed as requiring generalization to a larger environment and novel object compositions.
Regardless, it is interesting to consider a representation learning architecture that benefits each of these settings.

The main hypothesis, thus, is that a modular representation is desirable since individual modules can learn to capture these separate structures. This is achieved in CPS by having multiple different RNNs process the input separately and exchange information with one another to form a state representation, similar to how this is done in RIMs. Unlike in RIMs, however, here the input attention is over the features produced by an observation encoder, and attention is to all positions, while RIMs attend to particular spatial positions.
The motivation for this is not quite clear to me, other than that it is now possible to capture information about an object that covers multiple spatial positions. A potential downside is that it is no longer possible to attend to a particular location.

The technical contribution of CPS and technical novelty compared to prior work is therefore highly limited, especially since similar forms to the multiplicative input attention across positions explored here have also been explored previously, eg. in FiLM (as is also correctly pointed out in the related work section). Regarding the motivation of CPS based on "perceptual schemas", the connection is not quite clear to me either. In essence, the modularity in CPS and the reason why this may be beneficial for generalization is identical to how this is motivated in prior work such as RIMs. Therefore, it is not clear why modules or RIMs are now rebranded as perceptual subschema's, and one should view the learned representation as a perceptual schema. What exactly is new here that warrants a connection to "schemas" and a rebranding of existing terminology from previous papers?

The empirical contribution of CPS does appear substantial since in all the considered domains CPS is shown to outperform the baselines, which include RIMs. However, there are several experimental results regarding CPS and baselines that I would like to see clarified:

* On the Ballet task it is found that either feature attention or modules are sufficient for generalization. At the same time, RIMs and AAA are shown to perform very poorly. The only difference (minus minor architectural differences) of CPS with no attention to RIMs is the lack of a recurrent encoder in case of RIMs. Is the hypothesis thus that the lack of recurrent features is mainly responsible for these baselines performing poorly? If this is the case then it seems that this task is quite adversarial to allow for a fair comparison to RIMs. Would it be possible to endow RIMs with similar capabilities (eg. a recurrent input encoder) as a more representative baseline? Similarly, although AAA does have access to the recurrent input features it also performs poorly compared to the LSTM. What is the hypothesis for this behavior? Can it be tested?

* Similarly on the 3D Unity environment, the supposedly stronger baselines of AAA and RIMs struggle greatly compared to CPS and the LSTM. In this case it is not clear that recurrent input features are needed, and so the previous comparison to CPS with 1 slot + attention and to CPS with multiple slots and no attention would be particularly revealing here. I would like to see this comparison here or some other experiment to better explain the performance gap to CPS and AAA. 

The multi-level keybox experiments are more revealing since it is now demonstrated how there is a clear benefit to having attention and modules, which was not the case for Ballet and couldn't be observed on 3D Unity. At the same time, it can now also be observed how AAA performs quite well and the performance gap to this baseline is small. In general, although there is often a substantial gap to the baselines it is not quite clear to me yet which components are contributing the most, and therefore how significant the empirical contribution in reality is. The analysis of representations for regularly occurring events is interesting and indeed suggests some specialization among the modules as was also previously found in RIMs.

Finally, regarding the overall clarity and presentation of the work, I find that it could be improved substantially. In particular, the choice of different types of generalization seems rather arbitrary and it is unclear why these particular types are significant. Regarding the second and third generalization regime it is also not quite clear what the dimensions are that are being composed over. To me it seems that they should be the objects, yet the text talks about composable segments as well. The correspondence of perceptual schemas to parts of the learned representation is also a bit confusing. At first, it is suggested that perceptual schemas take on things like car motions, sedans, etc. yet these are the structures that are captured by the modules, which are referred to as subschemas. My question thus is, which are the schemas, the individual modules or their composition? Regarding the method details, I would strongly encourage the authors to include some pseudo-code for this algorithm, and better highlight what are the novel contributions. It also isn't clear how the policy network is structured to leverage the underlying modularity of the representation, which seems like an important choice to discuss. The text below eq. 4 about the potential capabilities of using attention in this way in combination with eq 3. was not clear to me, and it would be good to provide concrete examples of how the things outlined there could be achieved. 

Some minor comments:

* There is a mention of f_message, which is mentioned nowhere else. I assume this is an artifact of a previous version.

* I encountered several typo's, which a spell checker should reveal for you as well.

summary_of_the_review: This paper presents a method for learning modular representations for RL that allow for generalization in three regimes. The technical novelty is highly limited, but the empirical contribution is promising. At the same time, I would like to see a number of the presented results clarified and some additional analysis done to understand why the baselines perform so poorly and which parts of the proposed method contribute the most in different regimes.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The paper proposes a modular state representation learning architecture called Composable Perceptual Schemas to discover regularly recurring patterns in its observations into “schemas” and process them in a modular and factorized manner. The proposed model is tested for generalization in the deepRL setting along various axes such as varying object types, object numbers, distractors, task solution lengths etc. and shows promising improvements over other recurrent baseline models.

main_review: Strengths:

Interesting set of diagnostic tasks/environments have been developed to study generalization of deepRL agents along various axes such as object types, object numbers, sequence length of task solution etc. 

Experiments showing qualitative analysis of the subschema representations learned by the CPS model are interesting. More such qualitative studies/ablations or visualizations (various attention coefficients during an episode, subschema activity during an episode? etc.), could further help dissect and gain insight into the various components of the proposed CPS model and their roles in the generalization performance. 


Weaknesses:

Proposed model and architectural novelty is quite similar in nature to several of the follow-up works on RIMs [1, 2, 3, 4]. Discussion and comparison with highly relevant follow-up work on RIMs [1, 2, 3, 4] which also propose solutions to the same problems of i) how to share information among modules [1, 3] ii) learning type-specific update rules [2, 4] has not been addressed or evaluated. How does the proposed novel feature attention mechanism and the general CPS model architecture benefit over other solutions proposed by these works [1, 2, 3, 4] to the similar problems? 

Another issue is with the lack of relevant baseline models being used to judge the benefits of the proposed CPS model architecture and its novel feature attention mechanism. The RIMs baseline model struggles to learn on several tasks (Figure 3 (a) (b), Figure 4), sometimes shows even worse performance than the vanilla LSTM. RMC [5] which models relational information among memory slots and uses this information to update the memory slots using key-value attention would make for a good and (relatively) simple baseline model to compare against.  


Results:

It’s rather interesting that the vanilla LSTM is able to generalize better than RIMs on a few tasks such as Dancers (Fig. 3.a) and 3D Place X on Y task (Fig. 4). This is surprising since RIMs have inductive biases that explicitly favour modularity and compositionality in representation over the vanilla LSTM, which are beneficial to generalize on these tasks. Could the authors offer 
some explanation for this behavior?

Are the total number of learnable parameters always kept the same across CPS model variants with 1/2/4/8 subschemas in all the experiments? 

The paper would improve if the authors could add/move some of the visualizations from the Appendix to the section 4.3.1 to expand this discussion and qualitative analysis on the extent of specialization and modularity in representations/behavior captured by various subschemas. This would further strengthen the author’s hypothesis of “learning n subschemas, each with 1/n parameters of a single monolithic module….. encourages the discovery of regularly occurring structures.” 

I didn’t quite understand how to interpret the visualizations on the pairwise correlation between subschemas in Fig. 6.b. While they show that subschemas interact with one another during an event, shouldn’t this interaction be rather sparse? If the subschemas are to specialize to unique recurring input patterns in an episode, shouldn’t just one subschema ideally activate only for the recurring events (ex: pickup key, drop key etc.) it has specialized to represent? If so, then how can there be a high correlation between subschema activations over an event? Does that indicate that multiple subschemas are modelling the same recurring pattern/event amongst them? Could the authors provide some clarification on this?

Writing/Presentation: 

How are the attention coefficients (alpha^i) in the f_att function computed? 

$f_{att}^{(i)} (x_t, e_{t-1}^{(i))} = (Z_t … )$. What is $e_{t-1}^{(i)}$? Do you mean $c_{t-1}^{(i)}$?

“We used multihead-attention for $f_{message}$”. What is $f_{message}$? It was not defined earlier.

Please include highly relevant information about certain key agent/training/environment hyperparameters in the main text that is needed to judge the validity of empirical results/ablations. The reading experience can greatly improve if one doesn’t have to constantly jump from main text to Appendix to gather information needed to judge the experimental settings and validity of results.  

“... note that “Place A on C” and “Place B on D” are both in our training curriculum …” wording is confusing. Does this mean that the agent has seen the same task structure in training but not with the same objects to pick and place? 

Minor typos: 

“Reoccuring” -> “recurring”?

“Pickupable” -> “pickable”

“Policy-state” -> “state”



[1] Mittal et. al, “Learning to Combine Top-Down and Bottom-Up Signals in Recurrent Neural Networks with Attention over Modules”, ICML 2020.

[2] Goyal et. al, “Factorizing Declarative and Procedural Knowledge in Structured, Dynamical Environments”, ICLR 2021.

[3] Goyal et. al, “Coordination Among Neural Modules Through a Shared Global Workspace”,  

[4] Goyal et. al, “Neural Production Systems”, https://arxiv.org/abs/2103.01937

[5] Santoro, Adam, et al. "Relational recurrent neural networks.", NeurIPS 2018.

summary_of_the_review: The paper develops an interesting suite of diagnostic tasks/environments to study generalization in deepRL agents. While the technical contributions and results are interesting, it is difficult to fully judge the merits of the proposed model and the main contribution of the feature attention mechanism if it is not referenced, discussed or evaluated in the context of related prior works [1, 2, 3, 4] that study the similar problems and propose similar solution methods. Further, the paper would greatly improve with more qualitative analysis into the CPS model’s components and their role in its generalization performance. For these reasons I’m giving a score of 5. I’m willing to increase the score if my concerns are addressed.
review_title: null

Meta-review: This paper develops a mechanism for learning modular state representations in RL that organize recurring patterns into composable schemas. The approach combines modular RNNs as in RIMs (Goyal et al., 2020) with a dynamic feature attention mechanism. There were a variety of concerns in the initial reviews that were addressed by the authors through a set of clarifications and improved empirical analysis, substantially improving the paper. However, there still remain some issues in clarity of presentation and inconsistent empirical results, especially in the form of clear take-aways from the empirical analysis and broader insights from the paper, as detailed in the individual reviews. The authors are encouraged to take these aspects into consideration in revising their manuscript.
Decision: Reject