Spatially-Aware Transformers for Embodied Agents

We are grateful to Reviewer 44cr for their thoughtful critique and acknowledgment of the well-written nature of our manuscript and its contributions to spatial memory in transformer architectures. Your feedback has been invaluable in enhancing the quality of our paper.

Novelty of the Adaptive Memory Allocator (AMA)

We answered this in Common Response CR3.

Scalability with a Large Number of Places

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.

For retrieving the relevant memories, we think that it is shown in Exp-5. In the experiment, the model should generate the mini-patch on the facial image even though sometimes, the appearance looks similar (e.g., when the patch is near to the eye). Figure 7 shows that SAT-AMA can generate the mini-patches accurately based on the actions and spatial knowledge. This result suggests that SAT-AMA can effectively discern relevant memories through the spatial knowledge even when observation differences are subtle.

We thank the reviewer for giving this insightful comment. We will try to add this strategy as one of our baselines in our revision.

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.

[1] Le, Hung, et al. "Intrinsic Motivation via Surprise Memory." arXiv preprint arXiv:2308.04836 (2023).

[2] Saade, Alaa, et al. "Robust Exploration via Clustering-based Online Density Estimation." (2022).

We would like to extend our sincere thanks to Reviewer QvRD for their positive and encouraging feedback on our manuscript. We are grateful for the recognition of the importance of our research and the comprehensive nature of our work. Your constructive comments have been instrumental in refining our paper.

Concerns on Misleading Title

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

Concerns on Over-claiming

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. One good example is 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], which is essentially close to the proposed SAT-PM-PH model.

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.

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.

Missing Experiments

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.

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.

Following my previous point, I think this study emphasizes the scenario of memory constraint, but the experiments only considers small memory capacity, neglating the difference in representation (i.e., how to store an observation) and the difficulty in grounding from the queries to relevant memories. Overall, the generalization and practical influence of the proposed methods is not very clear to me.

We are grateful for your critical insight regarding our treatment of memory storage strategy and retrieval of relevant memories from queries.

We recognize that the memory storage strategy based on observation differences represents a valuable line of inquiry [1, 2]. As mentioned in our common responses, our primary research question in this paper is “what if spatial information is available as time indices to transformer-based episodic memory?”. While the memory storage strategy based on observation difference is indeed an important aspect, it is somewhat orthogonal to the main focus of our study. Nevertheless, we agree that such a strategy could be a robust alternative to FIFO memory and could complement our spatially-aware memory framework.

We express our gratitude to Reviewer XL1m for their insightful feedback and for recognizing the potential of our work to enhance episodic memory with spatial context. We are particularly appreciative of the acknowledgment of the significance of our proposed Adaptive Memory Allocator (AMA) and the comprehensiveness of our experimental validation.

Clarification on the Design and Implementation of Spatial Embeddings

Thank you for the constructive suggestion. We agree that our paper would become more clear and complete with the suggested improvement. As suggested, we made the following improvements.

In Section 2.1, we added the following.

“For the time and spatial embedding, the learnable embedding \citep{bert} or sinusoidal positional embedding \citep{transformer} could be applicable, but in this literature, for simplicity, the sinusoidal positional embedding is used.”

In Appendix C.2, we provided a comprehensive description of the spatial embedding construction for each experiment. For example, we clarified the fact that we used 2-D sinusoidal positional embedding in Exp-5 (the FFHQ generation experiment) while sinusoidal positional embeddings were primarily employed in other experiments.

In Appendix D.9, we added a comparative analysis of sinusoidal, 2-D sinusoidal (sinusoidal embedding for x and y axes), and 2-D learnable positional embeddings. This suggests that that 2-D sinusoidal and learnable embeddings are comparable while the sinusoidal embedding is clearly worse than them.

What’s the difficulty of learning AMA policy

We answered this in Common Response CR3.

Improving the Organization of the Paper

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.

Application to the more realistic embodied environments

Thank you for the suggestion. As suggested, we conducted additional experiment in a complex environment Habitat. See Common Response CR2 for more details.

Details of the experiments are missing. Input & Output of each environment, and training algorithms for different baselines

Thank you for the suggestion. As per your advice, we have included the suggested details of the experiments in Appendix C.1. These details include the exact input and output of each environment, as well as the representations that are fed into transformer memories.

Regarding the training algorithm of the baselines, we have provided the details in Appendix B.3. The training algorithm for SAT-FIFO is the same as SAT-AMA, except that it does not require learning to choose a strategy and therefore does not involve policy learning. In the case of DNC, we sampled a batch of episodes and directly optimized the downstream task loss function.

The detail models and hyperparameters of SAT variants and DNC are provided in Appendix B.4 and Appendix D.6, respectively. For reproducibility, we will also release the code upon the acceptance of the paper.

We express our gratitude to the reviewer for the recognition of our paper's motivation and the breadth of our experiments. As addressed in the common responses, we have discussed the extension of our work to the complex embodied AI experiment in the Habitat environment.

Usability in more complex embodied AI scenarios

As suggested, we added additional experiment results in a complex environment Habitat (suggested by reviewer XL1m). See Common Response CR2 for more details.

More discussion on AMA

We thank the reviewer for their insightful comments regarding the need for more in-depth discussion on the Adaptive Memory Allocator (AMA). We agree that providing a more comprehensive discussion on AMA would significantly enhance the paper.

As suggested, we investigated further to show the distribution of strategies chosen by AMA, and discuss its dependency on the nature of tasks, and the correlation between strategy choices and task performance. The table below presents the distribution of strategies selected by AMA, along with corresponding performance metrics, derived from the Ballet-FIFO task in Exp-3. We note that each column represents the probability of each strategy being selected, the accuracy when each strategy was selected. The last row represents the accuracies normalized for the four strategies.

The strategy distribution was averaged across three distinct runs (random seeds) and was ascertained by passing the output of the AMA's policy network through SoftMax. The results elucidate a clear correlation between the chosen strategies and task performance, reinforcing the premise that AMA's strategic choices are inherently linked to efficacy.

We added this new results in Appendix D.4 of the revised version.

Again, we appreciate your suggestion. We believe that by incorporating this analysis, the paper will become more comprehensive.

• To address the reviewer XL1m’s concern for clarifying the spatial embedding,
  • In Section 2.1, we added the detailed description for which embedding is used for the spatial embedding.
  • In Exp-5, we added the additional explanation of used spatial embedding for this experiment.
  • We added Appendix C.2 to give the detailed information of the utilized spatial embedding for each task.
• To address the reviewer XL1m’s concern for a lack of the detailed information of the experiments,
  • In Section 3, we added the additional explanation about the experimental setting.
  • We added Appendix C.1 to share more detailed information about the input and expected output for each task.
  • We updated Appendix B.3. to share the training algorithms of different baselines.
• To address the reviewer BEnq’s request for more discussion on AMA,
  • We added Appendix D.4 to investigate the correlation between the chosen strategies distribution and task performances.
• To address the reviewer QvRD’s request to investigate further for the learning-based read/write mechanism,
  • We updated Appendix D.6 to evaluate the Differentiable Neural Computer (DNC) [1] compared to our model for multi-task setting.
• To address the reviewers BEnq, XL1m, and QvRD’s concern for applicability and scalability in more realistic environment,
  • We added Appendix D.11 to evaluate the our model in the more realistic 3D environment, Habitat environment.