Paper: Spatially-Aware Transformers for Embodied Agents (ICLR-2024)

Reviewer comment (Theory):
The key idea for SAT seems to use separate networks for respective places. But can the architecture be useful in the case of a large number of places (i.e. what if $K -> \infty$ that results in almost zero size of memory for each place)?

Author response:
Regarding scalability, our Spatially-Aware Transformer (SAT) architecture is designed with parameter sharing across place memories to encode observations efficiently. By employing a common network for all place memories, the SAT model can effectively attend to and process information pertinent to each specific place. This architecture enables the transformer to selectively retrieve relevant memory in response to queries, regardless of the number of places. The number of place $K$ could not be infinite, the maximum number of place is the number of memory, in which, the Place Memory will be same to SAT memory without the hierarchy.

Reviewer comment (Presentation):
The title of this paper is very misleading. The paper does not propose any new formation of the transformer architecture to better model spatial information, but focuses on managing agents' episodic memory for addressing different downstream tasks.

Author response:
Your observation regarding the title's potential to mislead readers is well-noted. We added 'Memory' in the title.

Reviewer comment (Novelty):
This paper overclaims several contributions. The paper says "we are the first to motivate, conceptualize, and introduce the notion of transformers capable of utilizing explicit spatial information", but as the authors mentioned that defining and managing external memory have been extensively studied in previous machine learning literatures. It simply compares place- and time-centric store and hierarchical read methods that are intuitively suitable for different downstream tasks.

Author response:
We acknowledge the similarities between our SAT-PM-PH model and the topological map based memory used in the DUET agent for vision-language navigation. However, our work extends beyond the storage of keypoint observations by embedding the spatial knowledge within the transformer’s attention mechanism, which allows for a more nuanced and integrated processing of spatial information. For example, the place memory could be constructed per place, the each memory is encoded with more detailed spatial information. Other differences are that our model is designed to be retrieved without the spatial information and to optimize the efficacy on the retrieval through the hierarchical architecture, while the topological map based memory doesn’t. 

In light of your feedback, we will revise our manuscript to more accurately represent the state of the art and clearly delineate the innovative aspects of our work. We will adjust our language to avoid overstatements and ensure that our contributions are contextualized appropriately within the broader landscape of research in this area.

Reviewer comment (Novelty):
The use of topological graph that stores observations of keypoints for agents in visual navigation, e.g., the widely applied DUET agent [1] for vision-and-language navigation [2], is essentially close to the proposed SAT-PM-PH model.

Author response:
We acknowledge the similarities between our SAT-PM-PH model and the topological map based memory used in the DUET agent for vision-language navigation. However, our work extends beyond the storage of keypoint observations by embedding the spatial knowledge within the transformer’s attention mechanism, which allows for a more nuanced and integrated processing of spatial information. For example, the place memory could be constructed per place, the each memory is encoded with more detailed spatial information. Other differences are that our model is designed to be retrieved without the spatial information and to optimize the efficacy on the retrieval through the hierarchical architecture, while the topological map based memory doesn’t. 

In light of your feedback, we will revise our manuscript to more accurately represent the state of the art and clearly delineate the innovative aspects of our work. We will adjust our language to avoid overstatements and ensure that our contributions are contextualized appropriately within the broader landscape of research in this area.

Reviewer comment (Experiments):
The experiments of action-conditioned image generation (Exp-5) and reinforcement learning agents (3.3) are not very convincing. Exp-5 is not a practical setting and 3.3 is relatively simple, which cannot represent other reinforcement learning agents.

Author response:
> The experiments of action-conditioned image generation (Exp-5) and reinforcement learning agents (3.3) are not very convincing to me. Exp-5 is not a practical setting and 3.3 is relatively simple that cannot represent other reinforcement learning agents, see more below.

To handle this, we additionally evaluated our model in realistic 3D environment, Habitat environment (suggested by reviewer XL1m). See Common Response CR2 for more details.

Reviewer comment (Experiments):
This paper lacks investigation on the more recent embodied agents and their memory management approaches.

Author response:
> I think this paper lacks investigation on the more recent embodied agents and their memory management approaches.
> 

We thank the reviewer for this constructive feedback. We acknowledge that we overlooked some recent embodied agents such as DUET agent. We will update the experiments with the more recent embodied agents in our revision.

Reviewer comment (Experiments):
The experiments presented in this paper are relatively small-scale and simple. There is concern about how the proposed methods might impact recent research that often applies more capable networks (and massive data) to learn generic spatiotemporal priors to facilitate representing and memorizing the observations. For example, in Figure 3, an agent might be able to create a very compact representation for all ballet rooms from a single visit to each room.

Author response:
> The experiments presented in this paper are relatively small-scale and simple. I am concerning how the proposed methods might impact recent research that often apply more capable networks (and massive data) to learn generic spatiotemporal priors to facilitate representing and memorizing the observations. e.g., Figure 3 - an agent might be able to create a very compact representation for all ballet rooms from a single visit to each room?
> 

We appreciate your constructive feedback. As we replied for your concern above, we additionally evaluated our model in the more realistic 3D environment, Habitat environment, which is discussed in our common responses.

Reviewer comment (Experiments):
A valuable baseline is missing: without explicitly defining any memory management strategies, but giving the agent a certain memory budget and asking it to learn to update the memory by itself by training the agent on a mixture of data and tasks (e.g., a more general version of DNC mentioned in Appendix D.4).

Author response:
>I think a valuable baseline is missing here: without explicitly defining any memory management strategies but gives the agent a certain memory budget and asks it to learn to update the memory by itself by training the agent on a mixture of data and tasks (e.g., a more general version of DNC mentioned in Appendix D.4).

We appreciate your suggestion to evaluate the Differentiable Neural Computer (DNC) across a mix of tasks. Even though we evaluated DNC for a single task in the Room Ballet environment, but we extended this experiment to cover the setting with multiple tasks. During this process, we observed that the memory capacity allocated to the DNC in prior experiments was less than that of the SAT-AMA. To address this discrepancy, we reassessed the DNC performance on a single task with a memory capacity now comparable to that of SAT-AMA.

These results are detailed in **Appendix D.6**. In summary, when memory capacities are matched, the DNC demonstrates quicker learning than SAT-AMA in terms of training steps. However, it is notably less efficient in terms of wall-clock time, and the performance is unstable throughout learning process. The slow learning speed in terms of wall-clock time is attributed to the recurrent nature of the DNC's architecture, which, unlike the parallel training capabilities of the Transformer, necessitates a slower training process.

Furthermore, the AMA's policy requires a period of exploration during the initial stages of learning, which may contribute to a comparative delay in learning speed when measured by training steps. These insights and additional details are expounded upon in the Appendix. We are grateful for your constructive feedback.

Reviewer comment (Presentation):
The organization of the paper could use some improvements, such as the missing Table 1 in the paper though mentioned in Exp-2, the confusing combination of the figures from different experiments, and many details deferred to the appendix, making it difficult to understand the details of the model or the experiments.

Author response:
Thank you for the constructive suggestion. We agree with the potential confusions and need for update. Unfortunately, due to the page limit, we are currently facing difficulties in incorporating certain Figures and Tables from the Appendix into the main paper. We will, however, continue exploring options to appropriately arrange the content. For instance, if the paper is accepted, it would be much easier to make the suggested updates as we would be allowed an additional page. In the meantime, we have clearly indicated in the paper that Table 1 is located in the Appendix.