Spatially-Aware Transformers for Embodied Agents
Authors: Junmo Cho
Jaesik Yoon
Sungjin Ahn
Venue: ICLR 2024 spotlight (2024)
Keywords: Episodic Memory
Spatial Inference
Prediction
Generation
Reinforcement Learning
Primary area: general machine learning (i.e., none of the above)
TL;DR: We propose a transformer-based episodic memory model, the Spatially-Aware Episodic Transformer, that incorporates both temporal and spatial dimensions to improve memory utilization and downstream task accuracy.

Abstract: Episodic memory plays a crucial role in various cognitive processes, such as the ability to mentally recall past events. While cognitive science emphasizes the significance of spatial context in the formation and retrieval of episodic memory, the current primary approach to implementing episodic memory in AI systems is through transformers that store temporally ordered experiences, which overlooks the spatial dimension. As a result, it is unclear how the underlying structure could be extended to incorporate the spatial axis beyond temporal order alone and thereby what benefits can be obtained. To address this, this paper explores the use of Spatially-Aware Transformer models that incorporate spatial information. These models enable the creation of place-centric episodic memory that considers both temporal and spatial dimensions. Adopting this approach, we demonstrate that memory utilization efficiency can be improved, leading to enhanced accuracy in various place-centric downstream tasks. Additionally, we propose the Adaptive Memory Allocator, a memory management method based on reinforcement learning that aims to optimize efficiency of memory utilization. Our experiments demonstrate the advantages of our proposed model in various environments and across multiple downstream tasks, including prediction, generation, reasoning, and reinforcement learning. The source code for our models and experiments will be available at \href{https://github.com/spatially_aware_transformer}{https://github.com/spatially_aware_transformer}.

Published as a conference paper at ICLR 2024
SPATIALLY-AAAR TRANSFORMER FOR EMBODIED AGENTS
Junmo Cho'r, Jaesik Yoon 12,,, , Sungjin Ahn' KAIST& & SAP
ABSTRACT
Episodic memory plays crucial role in various cognitive processes, such the ability to mentally recall past events. While cognitive science emphasizes the significance of spatial context in the formation and retrieval of episodic memory, the current primary approach to implementing episodic memory in Al systems is through transformers that store temporally ordered experiences, which over- looks the spatial dimension. As a result, it is unclear how the underlying structure could be extended to incorporate the spatial axis beyond temporal order alone and thereby what benefits can be obtained. To address this, this paper explores the use of Spatially-Awee
Transformer models that incorporate spatial information. These models enable the creation of place-centric episodic memory that considers both temporal and spatial dimensions. Adopting this approach, we demonstrate that memory utilization efficiency be improved. leading to enhanced accuracy in various place-centric downstream tasks. Additionally, we propose the Adap- tive Memory Allocator, memory management method based on reinforcement learning that aims to optimize efficiency of memory utilization. Our experiments demonstrate the advantages of our proposed model in various environments and across multiple downstream tasks, including prediction, generation, reasoning, and reinforcement learning. The source code for models and experiments will be available at
-wwae aaaariilll-aaaarr---affrrmee transformer
INTRODUCTION
Episodic memory is the type of memory that involves storing and retrieving events episodes from an individual's life. This memory plays crucial role in various human cognitive processes, such as enabling people mentally travel back in time and remember past experiences like "what you did yesterday" (Perrin & Michaelian 2017 Suddendorf et al.) 22099.. It is also foundational in imagining the future and making effective decisions (Schacter al. (017). Therefore, AI systems aiming to mimic human-like cognitive abilities would benefit from incorporating a form of episodic memory.
Currently, the primary approach for modeling episodic memory AI is through transformers (Parisotto et al., 2020 Lampinen et al. 2021). This approach represents an episodic memory as a sequence of experience frames. During the storing process, frames are stored in First-In-Firss Out (FIFO) order with temporal-order embedding, called positional embsedding. While this FIFO inductive bias seems reasonable in modeling episodic memory due to the uni-directional flow of time in the physical world, various studies in cognitive science suggest that the formation and retrieval of episodic memories rely not only on the temporal aspects but also on another foundational axis of the physical world, the spatial axis, such where we parked car this morning (Buzsáki & Tingley 2018 Ekstrom & Ranganath] 2018 Pathman et al. 2018 van Asselen et al.) (006)..
Despite this significance, the spatial dimension has largely been overlooked in the modeling of transformer-based episodic memory. This can be attributed to two main reasons. First, since their invention, transformer models have mainly advanced in the language domain where spatial dimensions are not as prominent or meaningful. Second, in some domains, spatial annotations, such as the position of robot, may not be easily accessible as the temporal order indices. As result, previous research have concentrated mainly on improving the accuracy of inferring such spatial
"qqual contribution. Correspondence Sungjin Ahn esjn.ahn@gmail.com

---

Published as a conference paper at ICLR 2024
information from experience. e.g., sequences of observations and actions (Fraccaro et al. et al. 2020 Gregor et al., 2019 Whittington et al. 2022)..
2018
Uria
However, beyond language, there are broad range of domains, such robotics and virtual agents in games, where the spatial dimension is just as natural and crucial as the temporal dimension. These agents need to navigate around rooms and various areas, and having spatially structured knowledge such as "what happened where" is essential for achieving the goal. Moreover, in many cases, the spatial information is easily accessible as the time indices. In virtual worlds such games, for instance, the ground-truth position information is easily provided from the game engine (Guss et al. 2019) and for robots, it already popular to provide such spatial information by a separate hardware system such as BLE beacons (if it is in-door) (Faragher & Harle 2013), GPSs (EI-Rabbany) [2002). or
algorithmically, e-g., SLAM (Taketomi et al. 2017 7.aang et al. 2017 Despite these circumstances,
transformers currently only consider time axis.
In this work, we pose the question: what if spatial information is available as time indices to transformer-bused models? Our goal is to investigate this overlooked aspect of transformers and to examine whether and how transformers can utilize the spatial axis. We approach this by deviating from the traditional focus acquiring spatial information. Considering the widespread use of transformer architectures, as well as the growing accessibility of spatial information for embodied agents, we believe it is appropriate and timely address this question.
To achieve this goal, we propose a series of transformer architectures called Spatially-Awan Transformers (SAT or SA- Transformerr, which incorporate spatial information into transformers. This allows to explore various fea- sible design choices and assess the pros and cons of these architectures. Specifically, we demonstrate that Spatially- Aware Transformers offer three main benefits. Firstly, they enhance spatial reasoning tasks. We show that without explicit spatial information (only with temporal informa- tion), it is challenging to solve various general spatial reasoning tasks, such as "What has happened in the next room the agent?" (as shown in ( However, providing spatial information makes these tasks much eas- ier. Secondly, they enable efficient memory management, By incorporating spatial information, limitations of the FIFO memory, such as deleting impor- tant memory solely because it is old (as shown in Figure (b)). We can also implement structured approaches like place-centric hierarchical memory. We demonstrate that this structured memory improves effectiveness and effi- ciency in various place-centric tasks. To enhance memory utility further, we introduce the Adaptive Memory Allocator (AMA). Unlike traditional transformers that follows the FIFO policy for memory writing, AMA allows for adaptive selection of various writing strategies based on the goal. Finally, we demonstrate that these benefits extend across wide range of machine learning problems, including supervised prediction, image generation, and reinforcement learning.
Room A
Room
Room A
Room
Room
Room
(a)
(b)
Figure Home robot thought experiments. scenario (a), robot visits the environment in a different temporal order. When it arrives at Room D, answering what happened to the room its left would not be easy using only the temporal axis. scenario (b), the robot stays in Room B for a long time. With a FIFO memory. the agent would delete the memory of Room C. When it returns, it will perceive Room C as new.
overcome the
The main contributions of the paper are follows: First, to our knowledge, motivate, conceptualize, and introduce the notion of transformers capable of utilizing explicit spatial information. Second, propose series of SAT models designed to leverage spatial annotation. Third, we introduce the AMA method, which enables adaptable selection of various memory manage- ment strategies beyond the traditional FIFO-only strategy. Finally, we demonstrate that the advantages of SAT are not limited to a specific problem domain but extend across an array of machine learning applications, including supervised prediction, image generation, and reinforcement learning. We also release the source code for reproducibility.
the first to
PROPOSED MODELS
In this section, we describe the proposed Spatially Transformer (SAT) models. We begin with the simplest model, SAT-FIFO, and then discuss incorporating Place-Memory into SAT. Lastly,
2

---

Published as a conference paper at ICLR 2024
Drop
Tima
Place
Time
Time
Drop
M
ang max Qolr, - TAMAA
M
ang mie Qu(r(,)) -
(a) FIFO Memory
(b) Place Memory
(c) Adaptive Memory Allocator (AMA)
Figure 2: Illustrations of FIFO, Place Memory, and the Adaptive Memory Allocator (AMA). The orange memory is oldest and originates from a different place compared the green and blue ones. In FIFO, orange memory is dropped due to its age. However, the Place Memory system retains it through place- wise memory allocation. AMA chooses the appropriate allocation strategy. TAMA. based on the given downstream task knowledge, denoted as
introduce Adaptive Memory Allocation method for selecting a strategy other than FIFO based on the purpose of the downstream tasks.
2.1 DESIGN 1: SAT WITH FIFO MEMORY
When location annotation l, is available at each time step f, the most direct approach to make spatially-aware episodic memory is simply to introduce the spatial embedding e/oc embed(4,) and combine it with the time embedding (1 eféme and the observation embedding e For the time and spatial embedding, the learnable embedding (Devlin et al. 2018) or sinusoidal positional embedding (Vaswani et al. (017)) could be applicable. but in this literature, for simplicity, the sinusoidal positional embedding is used. This results in the experience frame x = sum_ embed,.(e (efoe, etime, egbe). Then, add the experience frame the episodic memory M, in the conventional FIFO order. We call this simplest model SAT-FIFO and its memory structure is illustrated in Figure
The main advantage of this approach is that it requires minimal alteration from the standard FIFO transformer in implementing the three basic operations of an episodic memory: WRITE (M++1 Write(M,, READ (ı/² Read Read(M, 9))). and REMOVE (M,+ Remove (Me, o)). Here, qo represents the query used to retrieve information from the episodic memory, and is the output of the retrieval operation that will be used to solve downstream tasks. When it comes to removal, we generalize the operation by introducing a strategy that determines how to remove experience frames. In the case of SAT-FIFO, the strategy a is fixed to FTFO. Although this approach endows the standard transformer with spatial awareness with minimal effort, it has limitations. For example, when the memory full, the SAT-FIFO model has to remove the oldest experience first even if it is important to keep it for achieving the goal of the downstream task, e.g., illustrated in Figure T (b).
2.2 DESIGN 2: SAT WITH PLACE MEMORY
To resolve the limitations of SAT-FIFO, we introduce the Place Memory (PM) into the Spatially- Aware Transformers. In this model, called SAT-PM, the whole space is represented by number of places and we maintain a place-wing wise episodic memory M one per place k = I,... K, resulting in place-centric hierarchical episodic memory M, {M+ A= shown in Figure 2b
place can be viewed as a cluster of locations and we observe that this place structure is naturally provided in many applications, e.g., via BLE beacons for indoor robots, via GPS and maps for autonomous driving cars, and via clustering algorithms in general. Even if the place structure must be learned through clustering algorithm, the extremely low dimensionality (only or 3) of the spatial coordinates. Moreover, as we shall show in Exp-5 in Sec. 3.2, is not always necessary have the ground truth place structure but an approximate one can be adequate in practice. Thus, in our experiments, we assume that the location is already clustered and thus replace the location embedding with place embedding place
note that is relatively simple learning task due to
Place-Centric Store: To store the experience frame N obtained at place a FIFO order. The length of each place's 's episodic memory, Lp, can be set to L/ by default or
added M, in
'Technically, it is not time embedding, but position (or temporal order) embedding. However, for simplicity, refer to time emhedding throughout the paper because "position" can also refer spatial position.
3

---

Published as a conference paper at ICLR 2024
manually adjusted with the prior's guidance. This allows the Place Memory retain the memory of a room, regardless of how long ago it was visited, as long as the place's ' memory is not full. Unlike the flat FIFO transformer, this approach provides greater flexibilityy, balancing total memory capacity between time and space and help mitigate the issue shown in scenario (b) in F However, if place's memory is full, experience frames will still be removed in a FIFO order. Although such deletions are inevitable in any episodic memory with finite capacity, in the next subsection, we will discuss different memory management strategies that can handle this situation better than FIFO.
Cenercen Hierarchical Read: A critical issue with transformers is their computation cost increasing with the number of memory items (Vaswani et al. 2017). One approach to mitigate this problem in an episodic transformer is to group set of experience frames (Lampinen et al.) ( into chunky and perform hierarchical read operations. Specifically, we first select the top-k chunks by applying attention only across the chunk representations instead of individual experience frames. Then, we search for experience frames within those chosen chunks, as shown in Figure 9
Previously, only time-based chunking was available. However, we observe that this approach can lead to problems in spatial environments. This because experience frames place can be scattered across chunks located at various positions in the flat memory (depending on the visit order of trajectory), and similarly, a chunk can contain mixed experiences from multiple places. This makes it more difficult to learn to search for these frames and develop representation of what happened in a place. Utilizing spatial awareness, our Spatially-Aware Transformer provides a simple solution to address this issue. We first ensure that a chunk is filled only with experiences of the same place, and these chunks are managed within a place memory. In experiments, we show that this simple place-centric hierarchical reading is both fast and effective for spatial reasoning. For more details about the hierarchical operation, please refer to the Appendix B.2
2.3 DESIGN 3: ADAPTIVE MEMORY ALLOCATOR
So far, we have assumed that experience frames are added to each place memory in the FIFO order. While using the inductive bias. "recent knowledge is more important", reasonable in modeling languages, it may not hold for various spatial problems. For example, a task may simply require retaining memory of a place that occurred at the beginning of an episode. Hence, to have a generic episodic memory, a more desired approach is to decide what to remember or forget adaptively according to the goal of the task.
However, a dilemma exists in practice. If we prioritize flexibility, the best approach would likely be to learn the optimal decision of which knowledge to be updated at each time step (Graves et al., 2014)) 1016|Hung et al. 2019). However, it is known that this approach is notoriously difficult to train in practice, and we show it empirically in Appendix D.6 In fact, transformers completely remove the problem of learning to write into the memory by adopting the FIFO policy. This is one of the most crucial advantages of the transformer model, which allows it to be scalable in the language domain but at the cost of losing adaptability to find better writing strategy.
To address this issue, we propose a simple and practical method called Adaptive Memory Allocator (AMA) that balances flexibility of the writing policy and ease of training. AMA extends the space of memory management strategies to multiple strategies by allowing the developer to design the space of various memory management strategies, denoted as A = _o, 7 }. based on their prior knowledge. For instance, can be FIFO, T2 can be Last-In-First-Out (LIFO), dg can be Seest-Visited-Firs-OOe (LVFO), and on. This way, the traditional transformer can be viewed as a special case where there is only one strategy, i.e., FIFO. Also, these strategies are similar to CPU-cache management algorithms (Denning 2005), but our goal to learn how to choose between them. It is important to note that the primary value of this method lies not in what the strategies themselves are, as they can be created various ways based on prior knowledge. Rather, the value lies in the fact that Spatially-Aware transformers enable the selection of these strategies (other than FIFO) in a way that optimizes downstream tasks.
The problem is then formulated as choosing the best strategy based the agent's current intention or task description т. For example, the agent can be given T = "What object did you see in your left room?" This problem can be addressed by learning policy (T) We implement this Q-learning policy TAMA = arg e Qss (T, a), where ø is the parameter of the value function. One might think that, if the number of strategies is small enough, we simply try all strategies and
one-step
4

---

Published as a conference paper at ICLR 2024
pick the best one. Then, we do not need to learn how to choose one. However, there are almost infinitely many ways of describing a task, e.g., via language, it should be able to generalize to the variations of task descriptions. By learning the context-aware policy in AMA, a suitable strategy even if an unseen (as a variation of) task description is given at test time. AMA is illustrated in Figure 2c| and the pseudo-code for learning algorithm with SAT and AMA is presented in the Appendix B.3
can learn to choose
EXPERIMENTS
This section presents an empirical evaluation of the Spatially-Aware Transformer (SAT) and Adap- tive Memory Allocator (AMA) across various environments and downstream tasks. We begin by evaluating the models on prediction tasks in the Room Ballet environment. Then, we demonstrate its capability to build an action-conde conditonnd world model and spatially-awar image generation model. Lastly, we showcase the use of SAT-AMA for training downstream reinforcement learning policy. Each baselines and tasks are explained each the experiment section. the experiments, observation, the time-step, and space are sequentially presented to the agent. Afterward, query is provided. For some experiments, the task type is introduced at the episode's commencement to select the proper memory allocation strategy. The agent is expected to integrate these inputs to solve the tasks. The more details are in Appendix C.T
3.1 SUPERVISED PREDICTION WITH ROOM BALLET
In this section, we use the Room Ballet environment, inspired by the Ballet task described in (Lampinen et al. 202T). Figure 3 illustrates the environ- ment, which consists of multiple rooms. By default, there are rooms, but this can be expanded to 16 and 25 in Exp-2. Each room represents a "place" and contains a dancer. In each episode, the agent is placed randomly one of the rooms and moves through the rooms via random walk. Upon entering a room, the agent observes complete dance performance consisting of 32 frames, performed by the dancer in that room. In each episode, a dancer type (i.e., the appearance of a dancer) and a dance type (i.c., the performance) are randomly and independently assigned to each room.
III
H
X
Agent Observation
Figure 3: The Room
Exp-1, Implicit derivation of spatial information in transformers
Ballet environment.
To verify that the transformers derive the spatial information, we designed spatial reasoning task on the Room Ballet environment, which we call the Next Ballet task. In this task, the agent visits 40 rooms (1280 steps) through random walk with the four actions. Afterward, when provided with query image showing a dancer's 's appearance, it is tasked to predict the type of dance that was performed in the next room of the queried dancer in clockwise order.
Although transformers have demonstrated impressive performance across various applications (Brown et al. 2020) Dosovitskiy et al., 2020; He al., 2022 Parisotto al.) 2020 Micheli et al.. 2022). it is unclear whether they can infer spatial information from navigation trajectories consisting of observations, actions, and time. To investigate this, we evaluate the following models: standard transformer only with temporal order information T-FIFO, a transformer with temporal order and action T-FIFO+A, and a transformer with spatial information SAT-FIFO. Also, to ensure that memory capacity does not affect the results, we set the memory size to be large enough to store all observations.
As shown in the Figure 4 (a), we can see that the transformer using only temporal order without actions (T-FIFO) completely fails on this spatial reasoning task, suggesting its inability to develop an internal representation of the space from the ordered sequence of observations, as expected in the Thought Experiment. When we additionally provide the action information (T-FIFO+A), it began to show a better performance but still significantly inadequate to accurately solve the spatial reasoning problem by inferring the dancer's 's room location. These results suggest that inferring spatial information from observations and actions is not trivial in practice, though it may not be totally impossible. On the contrary, explicitly incorporating spatial information, SAT successfully solves the task, demonstrating the importance of having explicit spatial information in spatial reasoning tasks.
Exp-2. Place-centric hierarchical reading Time-centric hierarchical reading
5

---

Published as a conference paper at ICLR 2024
Next Ballet
Short-Stay
Ballet- Multtlask
1.0
1.0
1.0
0.8
agcurgcy
0.6
0.6
0.4
0.4
0.4
0.5
1.0
1.5
2.0
0.0
Rooms
16 Rooms 25 Roomrs
0.00
0.5
1.0
1.5
2.0
2.5
Training Staps
Training Stops
165
SAFFFO
SATFFO
SAFFFOOTH
541itandom
SaT-0M6-Lang
(b)
[c)
Figure 4: (a) The performance comparison of transformer memories equipped with different embeddings on the spatial reasoning task. (b) The performance comparison of unstructured memory and temporally/ssatiaa
structured memory spatial reasoning task in large environment. (c) The performance comparison of the AMA model and predefined memory allocation strategies on a spatial reasoning task.
To demonstrate that place-centric chunking proposed in Section 2.2 is more beneficial for spatial task compared to traditional time-based chunking, we introduce new and more realistic task called "Short Stay" in the Ballet environment. In this task, unlike the previous one, the agent is not required to stay in room until it observes the entire dance. Instead, the agent exhibits more natural behavior by being able to move to a neighboring room at any time step. After making 1280 such moves, the agent is presented with query dancer image and tasked with predicting its dance type. To verify the efficacy of the place-centric hierarchical SAT, we compared three models: (i) SAT-FIFO, (ii) SAT-FIFO with time-centric hierarchical read (SAT-FIFO-TH) and (iii) SAT-PM with place-centric hierarchical read (SAT-PM-PH)) The accuracies are measured after 200K training steps, and the chunk size is set 32. More details about this task are provided in Appendix
As shown in Figure 4 (b), it is evident that SAT-PM-PH consistently outperforms SAT-FIFO-TH and SAT-FIFO in all scenarios. Furthermore, the number of rooms increases, the performance gap between SAT-FIFO-TH and SAT-PM-PH also widens. We observe that this is because SAT- FIFO-TH allows for temporally consecutive experiences from different places to be grouped together within a chunk. Similarly, experiences from the same place can be spread across multiple chunks. Consequently, gathering information about an event that occurred in a specific location by collecting scattered experience frames across multiple chunks remains a more challenging task for SAT-FIFO- TH, despite use of place encoding and computational speed enhancement due to the hierarchy. On the other hand, SAT-PM-PH easily overcomes this challenge since a chunk contains only experiences from the same place. When compared to the -on-hierarchical SAT-FIFO, SAT-PM-PH not only achieves better accuracy but also demonstrates faster inference speed, as in Table iin Appendix. This because SAT-PM-PP only attends to the memory in the top-k relevant chunks, while SAT-FIFO needs to attend to the entire memory.
Exp-3. Learning to select memory allocation strategy with AMA
In this section, we assume that the memory capacity is limited and demonstrate that with SAT-AMA, is possible to learn to find a proper memory management strategy that best fits the downstream task and that it is better than always using FIFO. We first intro- duce four strategies: First-In-First-uut (FIFO), Last-In-Firrt-OUu (LIFO), Most- Most-VVsitedd Wisited-Fiist-OOi (MVFO), and Least- Visited-First Out (LVFO). FIFO/LIFO replace the oldest/newest observation in the memory. MVFO/LVFO replace the oldest memory from the most/least frequently visited room. Note that whether these strate- gies are good or not is not the primary value of this work. Rather, they are examples of prior strategies that can be designed.
AIF
U1O
MYTO
19F0
AM4
0.999
0.145
0.325
0.734
0.094
0.8
0.090
0.652
0.192
NVVO
0.313
04
0:372
0.703
0.133
0.097
0.395
0.0
Figure 5: The accuracy of differ- ent task-strategy pairs.
Then, we introduce four different tasks, one for each strategy: Ballet-FIFO, Ballet-LIFO, Ballet- MVFO, and B...et-LVVFO. the tasks, which consist of a total of 9 rooms as shown in Figure 3 the agent changes rooms 17 times (for total of 576 steps) using random walk. This allows the agent to visit a room multiple times, potentially encountering multiple distinct dancers in a room. It is important to note that we limit the memory capacity 288 and thus half of the total episode steps must be chosen to be removed according to the management strategy. We note that the agent changes rooms after observing complete dance performance in those tasks.
In each task, the agent is asked to predict the dance type of a queried dancer image. The selection of the query dancer follows specific strategy, ensuring that it can only be answered correctly if the

---

Published as a conference paper at ICLR 2024
Generation on FFHO
Generation on Miniworld
3D Visual Match RL Task
0.10
1.50
0.38
Success Rates
0.8
1.25
0.05
0.6
1,00
0.75
0.04
0.4
0.50
0.02
0.25
0.00
1980
FRROHA
SAFIFO
Sar mahi)
0.00 co 035 0
1.0 1.5. 150 3.25 3:00
035 6.55 6.5 1.0 1.3 1.0 13355333
Taining Steps
lcs
Interaction Steps 1c7
SATHMP
SATAMA
SATPIP
(y)
Figure 6: (a) The generation performance comparison for the navigation task on facial images from the FFHQ dataset. For SAT-PM, the numbers in parentheses the number of the approximated places. (b) The generation performance comparison on the MiniWorld environment, which is 3D non-grid environment. (c) The performance on the Visual Match task on the MiniWorld environment.
correct storage strategy was used. For instance, in the Ballet-LIFO task, the query dancer is always selected from the oldest half of the episode. Therefore, if the FIFO strategy was used, the memory does not contain the oldest half, making it impossible to correctly answer the query. However, LIFO chosen, the memory always retains the necessary experience to answer the query. Similarly, in the Ballet-MVFO/LVV
tasks, the agent is asked to predict the most recently observed dances within each room and the dances in the most frequently visited room, respectively. Importantly, given a task, the agent does not know what the best strategy is but needs to learn to choose it via AMA. Note that we use place-centric hicrarchical read for selecting MVFO/LVFO, and time-centrie hierarchical read for selecting in FIFO/LIFO in AMA. More details about the tasks are provided in Appendix C.4
Tigure/3Jeppice the performance of different task strategy pairs. As anticipated, we observe strong performance in the diagonal terms of the matrix when the strategy aligns with the task. Conversely, when the strategy does not match (i.e., the off-diagonal terms), the performance deteriorates. Notably, in the 5th column, we present the performance of SAT-AMA. SAT-AMA achieves comparable performance to strategies specifically designed for each task (the diagonal terms) as it learns to discover suitable strategies for each task.
Exp-4. Generalization of SAT-AMA to multi-task descriptions
In Exp-3, we demonstrated that SAT-AMA can learn to select an appropriate strategy for a single task. However, in this specific scenario, it is possible to find the correct strategy by trying out all configurations,
there only four tasks and four strategies. In more realistic scenario, it is common to ask the agent to solve multiple tasks with diverse formats, such as language descriptions (as described in Sec. 2.3 In such cases, the model should be able generalize the diverse variations of task descriptions. In this setting, it is impossible to try all possible configurations.
To test this, we mapped each task in Exp-3 to a set of natural language descriptions. We created 20 different descriptions for each task. For example, for the Ballet-LVFO task, we provided descriptions such "Recall the dancers in the room the agent frequented the most" or "Reflect on the dancers in the room that seemed to be the agent's ss favorite". During training, we randomly selected 64 mixed descriptions (16 per task) and used the remaining 16 descriptions (4 per task) to evaluate All descriptions used in the
SAT with AMA. We name this version of SAT as SAT-AMA-Land experiment can be found in Appendix C.5 For AMA, the description given as an embedding from the OpenAl embeddings API The results are shown in Figure 4 (c). It is important note that the accuracy in the evaluation an average of three different vraining/evaluation description sets. this case, SAT-AMA achieved an accuracy of over 95%. This indicates that AMA can be extended to generalize to the diverse variations of task descriptions and has potential in broad applications such as language-instruce instructed robots.
3.2 EPISODIC IMAGINATION VIA IMAGE GENERATION
In this section, we demonstrate the potential of applying SAT to more complex tasks and environments: action- ceniionconditioned episodic image generation requiring multi-step spatial reasoning, and episodic image generation with AMA.
Exp-5. Action-condittooee generation in the FFHQ world
We took the embedding from the model textemmbedding-aaa-0022 /embeddingewhhtt-rr--emeedddddd
7

---

Published as a conference paper at ICLR 2024
FFHD Envionment
Acfion-based Denantiion Resul
GT
SAT-FIFO
600
T-FIFO
6.0
257
263
282
268
(8)
(b)
Figure 7: (a) FFHQ environment. The agent observes the partial view (yellow box) in the face image from FFHQ dataset. (b) Action-based generation result of SAT-FIFO and T-FIFO+A after Observation phase.
We demonstrate the capability of the SAT to generate conditioning on actions while navigating an environment to see its potential to implement an episodic world model (Gregor et al. 2019); Hafner ct al. 22200 Micheli et al.) 22022] Chen et al. (2022) To test this, we designed a navigation task on the facial images from FFHQ dataset (Karras et al. 2019), where 62,000 face images are used for training and 7, 7,000 images among the remaining images
used for evaluation. The environment is a face image which is considered as a 10 x 10 grid map. The agent observes a partial image of the face, as shown in Figure 7 (a), while navigating the image provided with an action (up, down, left, right) at each step. During the observation phase, the agent spawns in a random location and performs random walk for 256 steps. During the generation phase, a random position, it given sequence of random actions of length 32. Then, the task is to generate images corresponding to the queried action sequence by utilizing the episodic memory constructed during the observation phase. Therefore, this task requires multi-step spatial reasoning. See Figure 17 in Appendix for more visualization. For this task, 2-D sinusoidal positional embedding that encodes x and axes separately, is utilized for the space embedding. The comparative analysis of the embedding types for the space embedding is discussed in Appendix D.9
To verify the efficacy of the SAT-FIFO, we compared it with T-FIFO and T-FIFO+A. The results are shown in Figure 6 (a) and (b). Since this task is a harder version of the Next Ballet task in Exp-1 due to a larger map, T-FIFO and T-FIFO+A failed to solve it. SAT-FIFO, on the other hand, outperformed them and generated high-quality images. We note that this result is from the evaluation which contains facial images that are unseen during training. This result suggests that SAT can handle complex spatial reasoning on unseen visual environments. Additionally, we investigated the effect of using approximate place structures by changing the number of clusters (N 8, 16, 32). As shown in 6 (a), SAT-PM(N) showed robust performance across different place clustering. See Appendix D.7 for more results and discussion.
Exp-6. Image generation in 3D first-person view environment
Next, we evaluate SAT with AMA for the image generation in the MiniWorld environment (Chevalier Boisvert 2018), 3D environment where an agent navigates in the first-person view setting. As shown in Figure I1 in the Appendix, the agent placed at the center of a room with four dancers positioned at each side of the room (North, South, East, and West), each performing a distinct dance. The agent observes the room by rotating at a random degree between 3 and in a clockwise order. Afterwards, the agent is tasked with generating dance video featuring the queried dancer. Due to the rotation, during the observation phase, the agent observes the dance from various viewpoints, shown in Figure II] However, during the generation phase. the agent is asked to generate the dance from the front view. Therefore, the model cannot simply copy the dance from its episodic memory, but instead needs to develop proper representation of the dance from the episodic memory and decode accordingly. Similar to the Ballet tasks, we tested four different types of queries: FIFO, LIFO, MVFO, and LVFO. As depicted in F and [1] we can see that SAT-AMA successfully learns to choose the appropriate strategy for the given task. In contrast, SAT-FIFO performs worse because it loses necessary memory if it older than the memory capacity. More details about the experiment are provided in Appendix C.6
3.3 REINFORCEMENT LEARNING AGENT
Lastly, we test an RL agent equipped with SAT-AMA memory. The task involves two rooms. The agent is initially placed in room with single box of random color (either blue green). After staying four steps in the room, the agent is transported to the second room which is larger and contains multiple yellow boxes, and navigates for 80 steps.
8

---

Published as a conference paper at ICLR 2024
After this observation phase, the agent is warped back to the first room, where it is tasked to collect the box that matches the color of the box observed initially. The memory capacity is limited to 40, half of the number of total steps in the second room. We use Proximal Policy Optimization (PPO) (Schulman et al. (017)) for agent learning. AMA is designed to learn through the reward from the environment and to choose a strategy between FIFO and MVFO. As shown Figure (c), SAT-AMA successfully learned to select the appropriate strategy (MVFO) and solve the task. While it only needs to choose one out of two strategies, it is important to note that this still a challenging task where AMA needs to learn through sparse rewards and many interaction steps, and in the first-person view 3D environment.
Observation samples
Figure 8: RL environ- ment visualization.
RELATED WORKS
Transformers have been widely used in various applications (Brown et al. 2020 He et al. 2022
Parisotto et al.) 2020 Chen et al. 202T), however, their high computational cost has been major drawback, especially for reinforcement learning (Parisotto & Salakhutdinov [2021). To address this issue, several methods have been proposed to reduce the computational cost of transformers without sacrificing their long-term memorization ability (Parisotto & Salakbutdinoy [2021 Lampinen et al.,
2021]. In this paper, we propose a new type of hierarchical read operation that is more efficient for place-centric tasks. We also propose a memory management method based on reinforcement learning
that aims to optimize memory usage.
In parallel, there has been growing interest in using spatial information for embodied agents (Zhang jet al., [017 Rosenbaum et al. 2018: [Eslami et al. 2018 |iighh et al. 2019) Yoon et al. 2020).
Some studies have focused using spatial information to improve navigation performance (Zhang
cct al. al. 2018 Singh et al. 2019 Yoon al., 2020). In this study, investigate incorporating spatial
2017
Rosenbaum et al.
2018), while others have aimed to construct cognitive maps (Eslami
information into episodic memory. We show that this can improve performance and memory usage efficiency on spatial reasoning tasks without sacrificing performance.
CONCLUSION AND LIMITATIONS
Transformer-basee episodic memory has utilized temporal order to serialize experience frames. In this paper, we explore the potential of incorporating the spatial axis another fundamental aspect of the physical world. We argue that this crucial for embodied agents and introduce the concept of Spatially- from the simplest SAT-FIFO to SAT with place memory, which includes a hicrarchical episodic memory centered around places. Furthermore, introduce the Adaptive Memory Allocation (AMA) method, which provides a more flexible memory management strategy beyond the FIFO memory writing approach. Through experiments, we assess the improved performance of these methods and demonstrate their applicability to a wide range of machine learning problems.
Transformers (SAT). We propose different SAT architectures, starting
Limitations. Despite the aforementioned contributions, this work is based on several simplifying assumptions as the first to propose the concept of spatially-aware transformers. One of the main assumptions is its reliance on spatial annotation. Although spatial annotation is widely available in many embodied agent applications, incorporating the learning of spatial representation alongside its usage could enhance the applicability of the proposed models. Another limitation is the generality of the AMA method. While it offers more flexibility than FIFO, it still relies on set of predefined strategies. While it would ideal to discover the strategy from scratch, previous works like the Neural Turing Machine (Graves et al. 2014 2016) have demonstrated that this approach can introduce significant training complexity, limiting its practical application. Thus, there is a need for a new method that improves both flexibility and case of training. Lastly, while in this work focused on spatial reasoning tasks, there would also be non-spatial reasoning tasks in spatial environments. Generalizing the proposed method to handle such cases would also be an interesting future direction.
9

---

Published as a conference paper at ICLR 2024
ETHICS STATEMENT
In this paper, we present fundamental and general backbone for embodied agents using transformer that incorporates spatial information. Given the nature of this work, does not pose any immediate ethical risks it is primarily a theoretical contribution to the field of machine learning. However, we acknowledge the potential for this work to be applied in the development autonomous agents in the future. While such applications could have numerous benefits, they could also raise ethical concerns depending on the context and manner of use. We strongly advocate for the responsible and ethical use of this and any other Al technology.
REPRODUCIBILITY STATEMENT
To facilitate the reproducibility of our work, we have provided the pseudo codes, model architecture and hyperparameters in Appendix B.3 and B.4 We will also release the source code for models and experiments.
ACKNOWLEDGEMENT
This work is supported by Brain Pool Plus (BP+) Program (No. 202HH1D3A2203103355 through the National Research Foundation Korea (NRF) funded by the Ministry of Science and ICT.
REFERENCES
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam. Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
György Buzsáki and David Tingley. Space and time: the hippocampus as a sequence generator. Trends in cognilive sciences, 22(10):853-86, 2018
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.0948) 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 50844-50077 2021.
Maxime Bocovaaleer-Biiveerr Miniworld: Minimalistic 3d environment for rl & robotics research. http://gggttuub..oommaaiameeeeo
https://github. www.maximemmi.aaaammaa..ommm /yme miniworld] 2018.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint 19X1v:1901.02860. 2019.
Peter Denning. The locality principle. Communications of the ACM, 48(7):192 24, 2005.
Jacob Devlin, Ming-Wes Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXhv:1810.04805. 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer. Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 2020.
Arne D Ekstrom and Charan Ranganath. Space, time, and episodic memory: The hippocampus is all over the cognitive map. Hippocampus, 2809):680-687, 2018
Ahmed El-Rabbany. Introduction GPS: the global positioning system. Artech house, 2002.
10

---

Published as a conference paper at ICLR 2024
S. M. Ali Eslami, Danilo Jimenez Rezende, Frederic Besse, Fabio Viola, Ari S. Morcos, Marta Gar- nelo, Avraham Ruderman, Andrei A. Rusu, Ivo Danihelka, Karol Gregor, David P. Reichert, Lars Buesing, Theophane Weber, Oriol Vinyals, Dan Rosenbaum, Neil Rabinowitz, Helen King, Chloe Hillier, Matt Botvinick, Daan Wierstra, Koray Kavukeuoglu, and Demis Hassabis. Neural scene representation and rendering. Science, 36066399) 1210, 2018. doi: 10. 1tt.1///ciecc.arr11111 //wp::/wwwwwww. /tt./////o/tttt//////////////ora htt.a///ww/....aaooo...ooooo///a
aar6
Ramsey Faragher and Robert Harle. Location fingerprinting with bluetooth low energy beacons. IEFE journal Selected Areas in Communications, 33(1 (00111):218-2444 2015.
Marco Fraccaro, Danilo Rezende, Yori Zwols, Alexander Pritzel, SM Ali Eslami, and Fabio Vi- ola. Generative temporal models with spatial memory for partially observed environments. In International conference on machine learning, pp. 1549-1558. PMLR, 2018.
Alex Graves, Greg Wayne, and Ivo Danihelka. arXiv:1410.540,, 2014.
Neural turing machines. arXiv preprint
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using neural network with dynamic external memory, Nature, 538(7626): 471-476 2016.
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for rl. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.,
2019. URL https: http://:///ppcccccnnns...s www.:igg...uuuuurscc..uc.. faaeeerrffieeepppeeeee files//aperr22119 File/
William H Guss, Brandon Houghton, Nicholay Topin, Phillip Wang, Cayden Codel, Manuela Veloso, and Ruslan Salakhutdinov. Minerl: A large-scale dataset of minecraft demonstrations. arXiv preprint arXiv:1907.1344 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193. 2020.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 16000-16009, 2022.
Chia-Chun Hung, Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Carnevale, Arun Ahuja, and Greg Wayne. Optimizing agent behavior over long time scales by transporting value. Nature communications, 10(1):5223, 2019.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 4401-4410, 2019.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 1rXiv:1412.69888 2014.
Andrew Lampinen, Stephanie Chan, Andrea Banino, and Felix Hill. Towards mental time travel: hierarchical memory for reinforcement learning agents. Advances in Neural Information Processing Systems, 34:28182-22 28155 2021.
Andrew L Maas, Awni Y Hannun, Andrew Y Ng, et al. Rectifier nonlinearities improve neural network acoustic models. In Proc. ieml, volume 30, pp. 3. Atlanta, Georgia, USA, 2013.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. arXiv preprint 2209v:2209.00888 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human- control through deep reinforcement learning. nature, 518(7540:529-533. 533,, 2015.
11

---

Published as a conference paper at ICLR 2024
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th international conference on machine learning (ICML-10), pp. 807-814, 2010.
Emilio Parisotto and Ruslan Salakhutdinov. Efficient transformers in reinforcement learning using actor-learner distillation. arXiv preprini (tti::1004001555) 2021.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, al. Stabilizing transformers for reinforcement learning. In International conference on machine learning, pp. 7487-7498. PMLR, 2020.
Thanujeni Pathman, Christine Coughlin, and Simona Ghetti. Space and time in episodic memory: Effects of linearity and directionality memory for spatial location and temporal order in children and adults. PLoS One, 13(1 x3e022999,, 2018.
Denis Perrin and Kourken Michaelian. Memory as mental time travel. The Routledge handbook of philosophy of memory, pp. 228-239, 2017.
Dan Rosenbaum, Frederic Besse, Fabio Viola, Danilo J Rezende, and SM Eslami. Learning models for visual 3d localization with implicit mapping. arXiv preprint arXiv: 1807.03 /49, 2018.
Manolis Savva, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, Vladlen Koltun, Jitendra Malik, Devi Parikh, and Dhruv Batra. Habitat: Platform for Embodied AI Research. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2019.
Daniel L Schacter, Roland G Benoit, and Karl K Szpunar. Episodic future thinking: Mechanisms and functions. Current opinion in behavioral sciences, 17:41-50, 2017.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
Gautam Singh, Jaesik Yoon, Youngsung Son, and Sungjin Ahn. Sequential neural processes. Advances in Neural Information Processing Systems, 32, 2019.
Thomas Suddendorf, Donna Rose Addis, and Michael C Corballis. Mental time travel and the shaping of the human mind. Philosophical Transactions of the Royal Society B: Biological Sciences, 364 (1521):1317-1324, 2009.
Takafumi Taketomi, Hideaki Uchiyama, and Sei Ikeda. Visual slam algorithms: A survey from 2010 to 2016. IPSJ Transactions on Computer Vision and Applications, 91):1-11, 2017.
Benigno Uria, Borja Ibarz, Andrea Banino, Vinicius Zambaldi, Dharshan Kumaran, Demis Has- sabis, Caswell Barry, and Charles Blundell. The spatial memory pipeline: a model of egocen- tric to allocentric understanding mammalian brains. bioRxiv, 2020. doi: 10.1101/2020.11. 11.378141.
www. o bioooxiii vvrrxiv.oor org content/ ceatteent//artee ear wttpnttet/earr/2220/111111 //2p://///////////2222222
11.11. 378141
Marieke van Asselen, Rob Van der Lubbe, and Albert Postma. Are space and time automatically integrated in episodic memory? Memory, 14(2):232-24 2006.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
James C. R. Whittington, Joseph Warren, and Timothy E. J. Behrens. Relating transformers to models and neural representations of the hippocampal formation. 2022.
Jaesik Yoon, Gautam Singh, and Sungjin Ahn. Robustifying sequential neural processes. International Conference on Machine Learning, 2020.
Jingwei Zhang, Lei Tai, Ming Liu, Joschka Boedecker, and Wolfram Burgard. Neural slam: Learning to explore with external memory. arXiv preprint arXiv: 1706.0056691 2017.
12

---

Published as a conference paper at ICLR 2024
A COMPUTATIONAL OVERHEAD
Our study was performed on Intel server equipped with 8 NVIDIA RTX 3090 GPUs and 256GB of memory. The training duration for the tasks in the Room Ballet environment was approximately 12 hours single GPU. The generation and reinforcement learning tasks required approximately days and 12 hours on a single GPU, respectively.
B MODEL DETAILS
B.1 TRANSFORMER DESIGN
We utilized a memory-equipped transformer architecture, similar to the Transformer-X (TrXL) model (Dai al.) 2019). In this architecture, the query, key, and value not identical. Instead, the key and value are maintained as the equipped memory which is the same at every layer and is represented by the sum embedding of inputs such the observation or place information. The query is separate input, which is usually a single observation. This architecture avoids the quadratic computation required by naive transformers, but it does require attending to every memory. which can be burden in reinforcement learning tasks PParisotto & Salakhutdinov 2021
Guery
Time- Centric Hierarchical Read
Experience Frames
Quary
CenetrC Hierarchical Read
Figure 9: The agent randomly explores the places represented different colors, shown in the top left figure. The collected experience frames are shown the bottom left figure. When retrieving memory using time-centric hierarchical read, the chunks contain memories from multiple places, and memories from the same place are distributed across multiple chunks. This makes it difficult to retrieve place-centric knowledge from the chunks. In contrast, the place-centric read operation ensures that each chunk contains memories from the same place, making it easier retrieve place-centric knowledge than with the time-centric read operation.
B.2 HIERARCHICAL READ OPERATION
Our hierarchical read operation is inspired by the Hierarchical Chunk Attention Memory (HCAM) Lampinen al. (22T)) The memory is structured into chunks, which are sets of experience frames. The chunk size is a hyperparameter. In HCAM, each chunk consists of temporally contiguous memories, such as the memories from time step t+n. In our place-centric hierarchical read operation, each chunk is set of temporally contiguous memories that are associated with the same place, For example, if the agent visits place time steps 1, 1+2, and t+4, then those memories will form chunk, even though the memories at time steps t+1 and t+3 are not included. Each chunk has its own representative vector, which the average of the memories in the chunk.
The read operation is the same that in HCAM. At each layer, the query is attended to the chunk representatives, and the top-k most similar chunks are selected. The query is then attended to the memories in the top-k chunks, and the output of the read operation is represented by weighted average the chunk representatives and the memories the chunks, where the weights are determined by the similarities. By selecting the top-k chunks, the hierarchical read operation
13

---

Published as a conference paper at ICLR 2024
reduce the computational overhead of reading the memory. More detailed descriptions hierarchical attention mechanism is given as follows.
Hierarchical Attention Mechanism Now, explain how query extracts the relevant information from place memory, SAT with place-centric hierarchical read operation constructs a spatially-awann memory M {M ji=1- where M is a place-wise episodic memory per place i 1, ..... M. Each M' stores N; chunks, and we denote j-th chunk in M as Cij. Specifically, At time step t, when experience frame 24 in place is given, store the experience frame in the most recent chunk (i.e.. C.N,). ). If the chunk is full, then new empty chunk is allocated in M
By constructing place memory M from past observations, when normalized query g is given, it can use chunk representatives to choose top-k relevant chunks, and then attending within those most relevant chunks in further detail. In particular, let Caj. place denote the j-th chunk in M and place embedding at place i, respectively. Then the chunk relevance (Rig) computed as:
(1) where Linear(-) linear layer, pa-) is average over experience frames in chunk. Combine(:.-) can be any operation such as concatenation and addition, while we used addition in actual implementations. After retrieving chunks with top-k relevance scores, the model attends the detail within the chunks. Then, output query result q² is computed as:
Ray voftmax(Linear(g)) Linear(K)), K= Commbin((((Ci)) e((())
q"
Σ
Rg MMA Cijy Cij)
(2)
i.jElop-k from R.;
where MHA(-:.,)) is multi-head attention that takes query, key, value inputs in order. This attention mechanism lets query to select the most relevant memory chunks through space and time. For instance, query can extract information such as what happened in the left room of given context?.
Note that when the number of distinct place M is I, the formulation reduces to HCAM.
B.3 PSUEDO CODES
We introduce the pseudo-codes of SAT-AMA, SAT(-FIFO), DNC training for supervised learning task in Algorithm. 33
Algorithm 1 Spatially-Aware Transformer with Adaptive Memory Allocator
Require:
T: set of tasks with the task descriptions т. {oī H=) series of observations and location annotations. (g. y): query image and the true label. .). the downstream task loss. A: a set memory allocation strategies. e: a hyperparameter to explore. the parameters for
SAT and classifier. f.):: classifier for the supervised learing task. 6: the parameters for
AMA policy.
1: Initialize the parameters
2: while until converge do
3: Sample mini-batch of tasks T, from T
4:
for each task + E T, do
5:
With probability select a random rule aT, otherwise select a arg maX Qo(T,o) Initialize MT empty.
6:
7:
for t 1.... do
8:
/ = eemeed embble Memorize ME = Update(Mi
embed, em,ee embed.(o))))
9:
10:
end for end for
11:
12:
+- 3V. E TET, (Qo(T.o) C(fo(Read(MM. VVo ETEr, ([fo(Read((M
13:
8 - 0 - aoo
14: end while
B.4 MODEL ARCHITECTURES AND HYPERPARAMETERS
Model Architecture Our model architecture is based on the HCAM model Lampinen et al. (2021)
HCAM model is stack of memory layers, and each memory layer consists of:
14

---

Published as a conference paper at ICLR 2024
Algorithm 2 Spatially-Aware Transformer without AMA
Require: T: set of tasks with the task descriptions т. foz , ]=]==] series of observations and
location annotations. (g. v): query image and the true label. ((.,-): ): the downstream task loss.
0: the parameters for SAT and classifier. fo(.): a classifier for the supervised learing task.
1: Initialize the parameters
2: while until converge do
3: Sample mini-batch of tasks To from T
4: for each task r E Ty do
5: Select some strategy o" (e.g., FIFO) Initialize M3 as empty.
6:
7: for t , L do
8:
. = sum ses_ssueeeeeeseeeeeeeeeessssssse (embe(ee embedd)) embce embed. Memorize Mi = Up
9: end for
10:
end for
12: x- 0 - -- aVV.TET ((((Read((M))
13:
end while
Algorithm 3 DNC
Require: T: set of tasks with the task descriptions т. {or, 17 [= series of observations and location annotations. (g. y): query image and the true label. [(...): the downstream task loss.
0: the parameters for DNC and classifier. fo(.): a classifier for the supervised learing task.
1: Initialize the parameters
2: while until converge do
3: Sample a mini-batch of tasks from T
4: for each task т E Ty do
5: Initialize memory matrix Mj as empty.
6: for t = 1,..., L do
7:
27 = cat cat__aaeed__aaaeeeeeeaa embed, (embed())) (emeembedd.))eedd....e embed,(o7))
8:
Update MI = Write(M
end for 10: end for
9:
11: 4 0 - aVo ErET, (L
12:
end while
1. A local attention (LA) block, which uses self-attention within query vectors as in the transformer Vaswani et al. (2017).
2. A hierarchical chunk attention memory (HCAM) block, which allows the output query vectors from the LA block to perform cross-attention on hierarchical chunk memory.
3. A multilayer perceptron (MIP) block, which processes the final query output.
For each block, we use residual connection of the input, and layer normalization is applied before the block Parisotto et al. (2020). The transformer and SAT architectures are implemented by selecting the entire chunk at a high level with chunk size of 1. Different from HCAM, we use the same memory each layer, which represented by sum embedding as discussed in Appendix [B.]] Additionally, while HCAM stops the gradients through the memory, our model can be learned from the gradients through the memory.
Common Hyperparameters For all supervised learning tasks, we used batch size of 32, Adam optimizer Kingma & Bal (2014) with a learning rate of 0.0002, and gradient clipping over the range [-5.0, 5.0). We stacked four memory layers, each of which consists of an LA block, HCAM block, and an MLP block. We used a dimension blocks, used 2 heads and head dimension of 64 for multi-hese head attention operations. For the MLP block, we used a 2-layer MLP with a hidden dimension 128, and ReLU & Hinton (210)) activation function. For the Q network, used 1-layer ML. with hidden dimension of 64
128 for embedding vectors. For the LA and HCAM
15

---

Published as a conference paper at ICLR 2024
and ReLU activation function. We also used e-annealing for Q-learning Mnih et al. (2015). starting from 1.0 and decreasing to 0.2 over 200, 000 steps.
C TASK DETAILS
C.I INPUT AND EXPECTED OUTPUT FOR THE TASKS
As delineated in the main text, the agent acquires sequence of inputs observation, time-step and space information as the agent moves randomly except Exp-6. The agent movement in Exp-6 is discussed in Section C.6 Following these inputs, query is presented. For the Room Ballet and MiniWorld Environment, the query is an image showing a dancer's appearance. For image observations and query image, used CNN for the encoding. Then we added time and spatial embedding on encoded observations and feeding encoded observations and query to transformer as illustrated in Figure 9 For the conion--oniiionedd generation task, (Exp-5), the query consists of the observation at the starting position, and a sequence of actions for multi-step generation. Here, we used a learnable embedding for actions, and remainings are same. In reinforcement learning (RL) tasks, the query is the observation encountered after returning to the starting room. The encoding process is same in Room Ballet environment.
For the tasks where the memory capacity is limited such as Exp-3 and Exp-4, additional task information is imparted at the outset of the episodes, either numerical form or as narrative descriptions. The agent must employ its cognitive faculties to process and integrate these various inputs within the given constraints, such as limited memory capacity.
The agent's 's outputs are tailored to the nature of the task at hand: for the Room Ballet environment, it predicts the dance rype; for image generation tasks, it produces a sequence of images; and for RL tasks, it infers an action at each time-step.
C.2 PLACE EMBEDDING FOR THE TASKS
Room Ballet (Exp-1-4) The room formation is 3 to 5 5, and we used a ID sinusoidal embedding for place embedding. If room location is (3, 2), index 8 in sinusoidal embedding is used.
FFHQ Environment (Exp-5) The grid size 10 x 10. Here, we exploit 2D structure in place embedding. In particular, we posed disjoint sinusoidal positional embeddings for two sets, which 3, 4, 5, 6, 7, 8, 9, 10), and {11, 12, 13, 14, 15, 16, 17, 18, 19, 20} for x and y axis, respectively. For example, if the room position (x., v) = (2, 3). make the embedding the sum of sinusoidal positional embedding at 2 and 13. More detailed discussion is in Appendix D.9
MiniWorld Environment (Exp-6) We divided the map into 4 places, where each place covers 90 degrees range of agent's view (The agent can only rotate here). We used a 1D sinusoidal embedding.
RL Agent (Sec. 3.3) The environment consists of two rooms, and we used 1D sinusoidal embedding.
C.3 BALLET TASK WITH AN ACTION AT EACH TIME-STEP (SHORT-S STAY TASK)
The agent can randomly move one of four actions, up/downnleft/rigt at each time-step in this task. When it's movement blocked by the wall, then does not move. We tested x 3, 4 x 4, and 5 x rooms settings. The maximum movements are 1280, which is the same as the total observations in the Next Ballet task in Exp- The chunk size is 32, which is the same as the length of the dance to fit the configuration same other settings in Exp-1, 3 and 4.
C.4 BALLET TASKS IN EXP-3
As illustrated in Figure 10 in FFFOLL-FFFO/LIFO, the agent is asked to predict the dance type that was observed most recently in the past, respectively. In Ballet-MVFO/LVFO, the agent is asked to predict the most recently observed dances within each room and the dances in the most frequently visited room, respectively.
16

---

Published as a conference paper at ICLR 2024
Balle)-FIFFO Ask dance in second half
Ballet-LIFO: Ask dance in first half
Room
Ballet-LVFO: Ask dance in the most visited room
Ballet-MVFO: Ask the most recent dance among visited rooms
room, the agent observes dance while 16-step dalay before dancing (Total 32 steps)
Figure 10: The visualization of Ballet tasks in Exp-3 for 4-room case.
C.5 TASK DESCRIPTIONS FOR SAT-AMA GENERALIZATION EXPERIMENT
We utilized 20 task descriptions for each task ID for Q4. The task descriptions in Tables and
Secuance
Ground
AWN 19
stop 20
STRP 21
Step 22
DejeSt
SAT-AMA
5100 23
S109 24
stpp
301 27
2tp 25
stap 20
sbp
Figure 11: Left:
The overview of the Right: The generated samples from SAT-FIFO and SAT-AMA.
While SAT-FIFO failed generate dance correctly, SAT-AMA generated well even though saw the dancing while rotating.
C.6 GENERATION ON MINIWORLD
In this task, the aim is generate dance sequence based on observations of objects on North, West, East and South sides in the MiniWorld environment (Chevalier- 2018). The agent begins by turning around from North West side at a random degree between and 7, and stops to observe three objects dancing the South side. The agent then returns the initial angle and is asked to generate the dance sequence that was observed. The total number of steps is 126, with the agent turning 45 steps first, watching the dancing for 54 steps (during which three objects dance per 18 steps), and then turning again for 27 steps. The objects are chosen randomly from two shapes, ball and box, and six colors, including red, green, blue, purple, yellow, and grey. The dance sequences chosen randomly from five options: circle clockwise (move left move up move right move right move down move left), circle counterclockwise (move right move up- move left move left- move down move right), move up and down (move up move down move up move down move up move down), move left and right (move left move right move right move left move left move right), and move top-left and top-right (move tep- move bottom-right move top-right move bottom-left, move top-left, move bothom-right). The length of each dance sequence is 6.
The given task is one of FIFO/ B.4, except for the encoder and decoder. The encoder consists of three convolutional layers with hidden sizes of [64, 64, 16), kernel sizes of [8, 4, 3], and strides of [4, 2, 3]. The activation function used is LeakyReLU (Maas et al. 2013) with a slope of 0.01. After the convolutional layers, a single linear layer with hidden size of 128 used. The decoder architecture consists of one linear layer with hidden size of 144 and ReLU activation, followed by four transposed convolutional layers with hidden sizes of [64, 64, 18), kernel sizes of [3, 4, 8), strides of [3, 2, 4], and two ReL.U activations,
tasks and use common architecture in Appendix
17

---

while Published as a conference paper at ICLR 2024
well as single Tanh activation. It should be noted that the last layer's hidden size is derived from the image channel size multiplied by the length of the dance sequence, which is in this case.
D ADDITIONAL EXPERIMENTAL RESULTS
Inference Time (ms)
SAT
8.43 = 1.37
SAT-FIFO-TH
7.51 1.58
SAT-PM-PH
7.49
1.29
Table 1: The inference time comparison between SAT, SAT-FIFO-TH, and SAT-PM-PH. is the inference time only for SAT Read and the average value for 1000 samples. SAT requires attend the every memory, it requires much more time inference.
R1
R3
84
R5
96
R7
R8
Query
10
914
RB
10-
Attention
layer
10-2
Query
10-4
the type dancer in given query dancer?
Altention score at 4th layer
Figure 12: Left: The illustration the Opposite Ballet task. Right: attention scores at 3rd 4th layers.
D.1 ATTENTION MAP VISUALIZATION SAT
To verify that SAT can address the spatial reasoning task more clearly, we investigated the attention weight map visualization in SAT when solving spatial reasoning task. To do this, we introduced a new task in the Room Ballet environment called the Opposite Ballet task. This task is similar to the Next Ballet task in Exp-1, but there is no dancer at the center and the agent cannot go there. The agent observes every dance each room, and is asked to predict the dance that the dancer did the opposite room of the given dancer, as depicted in the left figure of Figure 12
We captured the attention weight map from SAT as shown in the right figure of Figure 12 We observed that the query attends to R2 (query dancer's room) in the 3rd layer (to extract the space knowledge in this attention), and then attends to the opposite room (R7) in the 4th layer. We note that the query has spatial information. This result implies that to do spatial reasoning, SAT requires several layers. SAT is implemented with multiple layers, it can infer the spatial reasoning even when the spatial information is not given in the query.
0.4
0.2
0.0
0.0
0.2
0.4
0.6
0.8
1.0
12
Training Steps
1e5
SAT-FIFO
SAT-PM
Figure 13: Accuracy on Ballet-ABA task. shows that SAT-PM handle the scenario (b) Figure SAT-FIFO fails.
18

---

Published as a conference paper at ICLR 2024
D.2 EMPIRICAL EVALUATION OF THE SCENARIO (B) IN THE FIGURE WITH SAT-PM
We discussed the Place Memory can handle the long-stap issue in the scenario (b) in the Figure || in the Section 2.2 Here, we empirically evaluated it the Room Ballet environment. To do this, we created new task, Ballet-ABA task. In this task, the agent visits total of 18 rooms (576 steps), with the capacity to store observations from only rooms (256 steps). The agent initially visits unique rooms, and subsequently remains in the last room for the duration of the task. The agent is then asked to predict the dance type of the given query dancer image in the initial rooms.
In this scenario, SAT-FIFO struggles to retain the necessary information for accurate predictions. In contrast, SAT-PM demonstrates superior performance, as shown in Figure [3 This underscores the advantage of utilizing aware memory storage operations when faced with memory limitations in spatial tasks.
0.8
0.2
0.0.0
0.5
1.0
1.5
2.0
2.5
3.0
6
10
Training Steps
le5
Wall Clock Time (hour)
SAT-FIFO with unbounded cache (Memory Size: 576 steps)
SAT-AMA (Memoryy Size: 288 steps)
Figure 14: Performance comparison between SAT-FIFO with unlimited memory and SAT-AMA on Ballet Strategy tasks. In this task, one of four tasks FBallc-FIFOOLIFOMM WWO./////////OOOOO////O is given for each episode as in Exp-3. SAT-AMA has to learn to choose the proper memory allocation strategy for the given task ID), while SAT-FIFO with unlimited memory requires retrieving larger memory.
D.3 PERFORMANCE COMPARISON BETWEEN SAT-AMA AND SAT-FIFO WITH UNLIMITED MEMORY
In Exp-3, we modified the problem setting to limited memory to show efficient memory usage of AMA. This made SAT-FIFO with limited memory to deteriorate in Ballet-Strategy tasks. One may ask that SAT-FIFO with unlimited memory can still solve the problem, however, we show that SAT-AMA outperforms SAT-FIFO in speed even SAT-FIFO has unlimited memory capacity.
To verify this, we evaluated SAT-AMA and SAT-FIFO with unlimited memory on Ballet Strategy that, even if it requires additional learning for the AMA
tasks in Exp-3. In Figure 14 policy, SAET-AMA can learn the task faster because it maintains small amount of memory. This result shows that learning through comparison with more memory slots is less efficient than learning with limited memory.
Strategy
FIFO
LIFO
MVFO
LVFO
AMA Distribution
0.475
0.073
0.150
0.302
Accuracy
0.999
0.146
0.325
0.714
Normalized Accuracy
0.457
0.067
0.149
0.327
Table 2: The relationship between AMA strategy distribution and task performance for the Ballet-FIFO task.
D.4 THE CORRELATION BETWEEN STRATEGY CHOICES AND TASK PERFORMANCES
To more deeply understand how Adaptive Memory Allocator (AMA) works, in this section, we illustrate how AMA adapts its strategy selection to different experimental settings. To address this, investigate 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 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
19

---

50 60 70 80 Published as a conference paper at ICLR 2024
probability of each strategy being selected, the accuracy when each strategy 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.
D.5 GENERALIZATION PERFORMANCE FOR THE UNSEEN SIZE OF ENVIRONMENT
To ascertain the model's generalization capabilities, we assessed performance variations under con- ditions where the environment size or memory capacity differed from the training phase. Our SAT model employs hierarchical memory organized in chunks (where each chunk is set of observations) and is designed to retrieve the most pertinent top-K K chunk memories in response sequently, the model is adaptable to changes in memory size, reflected by the number of chunks, ensuring consistent application.
query. Con-
In the MaIO task, detailed in Exp-3, the agent is trained in an environment with 8 rooms, visiting each room 18 times to observe the dancers' performances (remaining in each room for 32 steps and viewing the dance for half of the duration). The agent's 's memory is sized to accommodate observations from 9 rooms, enabling it to recall the top- (K=4) relevant chunks chunks-equivalers to observations from 4 rooms based on the query provided, infer the dance sequence.
To validate the model's adaptability to both environmental and memory size changes, we conducted supplementary experiments by varying the number of rooms to 4, 6, 10, and 12, using three random seeds for each scenario, and adjusted the memory size to match the number of rooms.
As Table 3 illustrates, the model demonstrated marginal increase in accuracy in environments with fewer rooms than the training setup, indicating an effective generalization. Moreover, the performance remained robust even with increased number of rooms, showcasing the model's resilience environmental scale variations.
8 Rooms (Train Env.)
Rooms
6 Rooms
Rooms
12 Rooms
Accuracy
99.54%
99.86%
99.76%
99.19%
98.04%
Table 3: Generalization performance of our model in environments of different sizes compared to the training environment.
1.0
0.8
Accuracy
0.6
0.4
0.2
0.0.0
0.2
0.4
0.6
0.8
1.0
1.2
0
10 20
30
Training Steps
le5
Wall Clock Time (hour)
FIFO
SAET
DNC
Performance comparison between DNC and SAT-AMA/FIFO Ballet-MVFO task. SAT-AMA, while DNC converges faster than SAT-AMA, the performance is worse, and is much slower than SAT-AMA.
Figure 15:
D.6 PERFORMANCE COMPARISON WITH DIFFERENTIABLE NEURAL COMPUTER (DNC)
One of the limitations our work is that we learn AMA policy from pre-defined set of strategies. In ideal case, the model should discover the strategy from scratch with learning-basse read/write mechanisms such as Neural Turing Machine (NTM) (Graves et al. 2014). However, is known that maaming-based read/write memory is slow and difficult to optimize practice.
To investigate the learning-base read/write memory, we applied Differentiable Neural Computer (DNC) (Graves al. 2016) some of our tasks, which is upgraded version of NTM. For this
20

---

Published as a conference paper at ICLR 2024
1.0
0.8
0.6
0.4
0.2
0.0.
0.5
1.0
1.5
2,0
10 30 40 50 60 70 80
Training Steps
1e5
Wall Clock Time (hour)
SAF-FIFO
AMAT-AMA
DNC firssirst
DNC.every
Figure 16: Performance comparison between DNC and SAT-AMA/FIFO on Ballet-MultiTask in Exp-3. DNC- first takes task ID as learnable embedding only at first time step. while DNC-every takes task ID as one-hot vector concatenated to an encoding of the observed image for every time step. While DNC-every converges faster than SAT-AMA, the performance is worse, and has much lower speed than SAT-AMA. DNC-first which has the same task setting as SAT-AMA (i.e., giving the task information only at initial step) performed worst.
investigation, it is evaluated for Ballet-MVFO and Ballet-MultiTsk in Exp-3. To apply this to DNC, The observation encoding is same to our model. Originally, DNC is not designed for the multi-tasks setting, we added the task information in the two ways. One of them, DNC-first giving the task information at the beginning of the episode as done for our model. Another, DNC-every is concatenating the task information to the observation embedding at every timestep. It is designed to prevent that the DNC forgets the task information because the backbone of DNC is the recurrent module which is limited to handle the long-term dependency. The DNC's memory size same to SAT-AMA, 288 cells with 128 dimensions each. LSTM was used for DNC's controller, and the sizes of image representation and representation in the hidden layer were set to 128.
The results are shown in Figures 13 and 16 Commonly, DNC showed more efficient training in the aspect of the training steps, while it is tremendously slower in the wall-clock time comparison. It is because the backbone of DNC is consisted of the recurrent module, which training is slower than the training of Transformer that is the backbone of SAT-AMA. The reason why SAT-AMA is slower in the training steps is that the policy of AMA should requires some exploration the early stage of the training. Additionally, DNC-every outperforms DNC-first, and it shows sub-optimal performance. This result suggests that DNC is limited to handle the long-term dependent knowledge. While the result on DNC models may require more training steps for better analysis. the result shows the limitations of the current learning-based read/write memory module. However, as stated in Limitations, investgating more advanced forms of learnable memory as strategy allocator would be an interesting future work.
D.7 SAT-PM WITH APPROXIMATE PLACE CLUSTERS IN FFHQ ENVIRONMENT
We investigate the effect of the approximate place structure on SAT-PM in FFHQ environment. In Room Ballet tasks, a place was already clustered rooms, and the agent could easily manage the place memory by storing the observations in each room separately. However, in FFHQ environment, there no pre-defined place structures. In realistic scenario, the agent should construct its own place clusters based on given locations with the clustering algorithm (e.g., k-nearest neighbor), then use the place structure for the downstream tasks.
To show that SAT-PM is robust to approximate place clusters, we designed three types of place clusters in FFHQ environment as depicted in Figure 17 (b). To generate place clusters, we applied a k-nearest neighbor algorithm with different number of clusters (N = 8, 16, 32) on set of locations that is available the environment (i.e., 100 locations in 10 x 10 grid map). Depending on the number of place clusters, we denote the SAT-PM model using those place clusters as SAT-PM(N). As described in Exp-6, SAT-PM(8), SAT-PM(16), and SAT-PM(32) showed comparable performance as SAT-FIFO. This implies that do not need the ground truth place structure but an approximate structure would be enough in practice. In Fig. (c), we present the full results of generated samples. We note that SAT-PM with approximate place clusters perform well as SAT-FIFO.
21

---

Published as a conference paper at ICLR 2024
(a)
(b)
FFHQ Environment
Phase
Examples
8 clusters
16 clusters
32 clusters
(c)
Action-based Generation Result
GT
SAT-PM(8)
SAT-PM(16)
SAT-PM(S2)
SAT-FIFO
T.FIFO+A
T-FIFO
257
263
282
288
Figure 17: (a) FFHQ environment. The agent freely navigates in 10 10 grid map observing partial (yellow box). In Observation phase, agent randomly navigates 256 steps (only the first half visualized in green trajectory). In Generation phase, the agent should generate the scenes given starting location's 's scene and action sequence (red trajectory). Circled locations are agent's 's initial location in each phase. (b) Place clusters examples. We use three types of place clusters with different numher of clusters using knn algorithm on set of available locations in grid map. (c) The generation result of all models.
D.8 SAT-AMA WITH APPROXIMATE PLACE CLUSTERS FFHQ ENVIRONMENT
As Exp-3 in Room Ballet environment, we can extend the FFHQ environment to limited memory setting where the choosing the proper strategy based on downstream task becomes crucial. To do this, we double the length of Observation phase to 512 steps, and assume that only half of total steps (i.e., 256 steps) can be stored in memory. In Observation phase, the agent stays in x 2 region for 256 steps in the middle of random walk. This scenario is from thought experiment in Figure (b)) where FIFO memory would fail to store full information of the facial image. However, AMA policy can choose place-centric memory which enables store full information of the image.
For evaluation, we compared SAT-FIFO and SAT-AMA. When place-centril centric strategy (MVFO) is chosen by SAT-AMA, used approximate place clusters (N 8) in Sec. D.7, In Figure 18 (a), we note that SAT-AMA with approximate place clusters outperform SAT-FIFO by large margin. In Figure (b), the visualization of generation results show that SAT-AMA could generate the scenes well, while SAT-FIFO fails to. This implies that SAT-AMA with approximate place clusters still work action-conditioned generation task.
D.9 COMPARATIVE ANALYSIS OF EMBEDDING TYPES FOR SPACE EMBEDDING FOR THE GENERATION TASK FFHQ ENVIRONMENT
In this paper, utilized the sinusoidal positional embedding Vaswani al., 2017) for simplicity. However, the space is not uni-directional. For example, it can be represented as x and coordinates. Therefore, investigated 2 dimensional positional embedding in this section.
We implemented two 2-D positional embedding; 2-D sinusoidal and learnable positional embedding. The 2-D sinusoidal positional embedding is constructed by summing sinusoidal embeddings corre-
22

---

Published as a conference paper at ICLR 2024
(a)
MSE
0.5
SAT-FIFO
SAT-MM(8)
(b) Generation examples (1st row: GT, 2nd row: SAT-AMA, 3rd row: SAT-FIFO)
Sample
Sample
Sample
Figure 18: (a) The generation performance on FFHQ enviromment when memory is limited. SAT-AMA outperforms SAT-FIFO large margin. (b) Generation examples of SAT-AMA, and SAT-FIFO.
Generation FFHQ
0.025
0.020
MSE Loss
0.015
0.010
0.005
0.000
0.0
0.2
0.4
0.6
1.0
Training Steps
1e6
Sinusoidal
2-D Learnable
2-D Sinusoidal
Figure 19: MSE loss comparison sinusoidal, 2-D sinusoidal, 2-D learnable embeddings for spatial embedding. The task pertains to image generation the FFHQ environment, as detailed in Section 33.2 The figure demonstrates the superior performance of the 2-D learnable embedding over the 2-D sinusoidal, with both significantly outperforming the standard sinusoidal embedding. This improvement attributed the non-unidirectional nature spatial dimensions.
sponding each axis. For instance, position (x.y)= (0,1), combine the Oth index sinusoidal embedding for the x-axis with the 1st index embedding for the y-axis. Notably, this method produces non-symmetric embeddings for symmetric positions, such (0,1) and (1,0), due the addition of large number to the index at the y-axis to ensure differentiation.
The 2-D learnable embedding follows a similar structure but is optimized through backpropagation based on the downstream task's 's loss. It is designed with distinct embeddings for the and axes, enabling discrimination of symmetric positions.
23

---

Published as a conference paper at ICLR 2024
As illustrated in Figure 19 our results indicate clear advantage of 2-D embeddings sinusoidal positional embedding (Vaswani et al.] 2017). with the learnable 2-D variant exhibit- ing marginally better performance than the sinusoidal 2-D. These findings suggest that spatial embedding approach that more accurately reflects the bidirectional nature of space is preferable. Interestingly, the additional flexibility afforded by learnability does not yield significant benefit, implying that the inherent structure of 2-D embeddings may be sufficient to capture the spatial relationships pertinent to the task.
the
1.0
0..8
0.6
Success
0.4
0.2
0.0.0
0.2
0.4
0.6
0.0
1.0
Room A (Step 1-4)
Room (Step 5-84)
Room A (Step 85-100)
Interaction Steps
1e7
SATT FIFO
SAT AMA
FIFO Memory
Place Memory
the goal ball color, and after steps in room B, the agent to collect the ball with same color to the ball room On the map, the red triangle is the agent, and light gray shows the area where agent can see. Top Right: Performance on Visual Match task. Bottom: The memories in FIFO and Place memory architectures.
Figure 20: Top Left: The Visual Match Task Visualization. In room A, the agent
D.10 VISUAL MATCH TASK ON 2D MINIGRID
We evaluated SAT-AMA for the Visual Match task (Hung et al. 2019). It is similar to the task in the section 3.3 as illustrated in the top left figure of Figure 20 but this environment, the agent should learn how to act from at the beginning of the episode where the agent observes one ball. As in the section 3.3 the memory capacity is half of the number of steps the second room, and PPO is used
also. AMA design is also the same. The result is shown in the top right figure of Figure 20 As in the
section 3.3 the SAT-AMA solves this task also, because the place memory can keep the memory at
the beginning the episode by managing the FIFO memories per place shown in the bottom figure
of Figure 20
D.11 SAT-AMA IN HABITAT ENVIRONMENT
To further investigate the applicability of SAT-AMA for embodied agents in realistic 3ID environment, we evaluated our model in Habitat environment (Sayva et al.] (019). Habitat environment is a 3D) simulator that allows the embodied agent to freely navigate house-like environment. We designed two tasks to demonstrate the applicability of SAT-AMA in this environment. The first task is classification task, which requires the model extract relevant patterns from complex scenes in order to perform well. The second task is generation task, which requires the model to generate unseen scenes based on spatially- aware memory.
For classification task, we designed the task similar to RL task in Section 3.3 and Visual Match task Appendix D.10 the beginning of the episode, the agent is randomly spawned in the house, and the agent can observe the small colored square (one of red, blue, green, cyan, grey, yellow) randomly located the scene (Figure c)). Then, the square disappears and the agent randomly navigates in the house. After navigating 200 steps, the agent stops at certain location and then randomly wander around nearby for another 200 steps. Thus, the agent observed total 400 frames in total. To this end, the first observation without the square is given of the square. To acquire the place information, we calculated the distance between the given location and center location of each room, and chose its place the closest room (Figure 21 a)).
query, and the model is asked to predict the color
24

---

Published as a conference paper at ICLR 2024
(a) Top-down view and trajectory
(b) Learning curve
1.0
0.8
Accuracy
0.6
0.4
0.2
0.0 0.0
0.5
1.0 1.5 2,5 3.0
Training Steps
1e4
SAT-FIFO
SAT-AMA
(c) Sample trajectory
Colored square
Figure 21: (a) Top-down view environment and agent's trajectory. Same colored points indicate same place. (b) Performance of SAT-FIFO and SAT-AMA on Habitat classification task. SAT-AMA learns to choose optimal strategy and attends to relevant memory in visually complex environment. (c) Sample trajectory of the episode. The agent observes the colored square at first step, then randomly navigates the house for 400 steps.
(a) Generation result
(b) Learning curve
2.0
1.6
1.2
SAT-
MSE
0.8
AMA
SAT-
0.0
2
3
4
5
Training Steps
1e5
FIFO
Figure 22: (a) Generation result SAT-AMA and SAT-FIFO. Random unseen direction is given as query, and SAT-AMA could generate the scenes well, while SAT-FIFO fails to. (b) MSE loss of SAT-FIFO and SAT-AMA.
To evaluate SAT-AMA and SAT-FIFO, we confined the memory capacity to 200. As illustrated in fully solve the task while SAT-FIFO performs near to random choice, about 18.75% accuracy. SAT-FIFO's failure is because of the limitation of the memory capacity (it only can store the last 200 observations). Whereas, SAT-AMA can memorize the first observation by choosing MVFO strategy. This shows that SAT-AMA can extract the relevant patterns in complex scenes.
Figure 2)),
For generation task, we designed the task inspired by (Rosenbaum et al. 2018 Eslami et al. 2018) At the beginning of the episode, the agent randomly spawned in the house, and the agent rotates to random direction repeatedly. After rotating for 150 steps (single rotation changes the direction by degrees). the agent randomly rotates for another 150 steps within a range 36 degrees in either direction from its current orientation. At the end, random unseen direction is given and the model is asked to generate the scene in given query direction. We use fixed set of unseen directions for simplicity. To acquire the place information, equally divided 360 into 10 regions, and assumed each region as place.
query,
To evaluate SAT-AMA and SAT-FIFO, we limited the memory capacity to 150. As illustrated in Figure (2(b), b). SAT-AMA outperforms SAT-FIFO in reconstruction loss. This is because SAT-FIFO stores last 150 observations which cover only some portion of the entire 360 scene. However, SAT-AMA can choose the MVFO strategy to store entire 360 scene effectively attending to relevant scenes for generation. The visualizations in Figure 22 a) shows that SAT-AMA could generate the
25

---

Published as a conference paper at ICLR 2024
scenes well, while SAT-FIFO fails to. This implies that SAT-AMA generate unseen scenes based on apat memory.
Tasks
Descriptions
Recall the first performers out of the 18 provided. Think back to the initial half of the 18 dancers. Keep the first nine from the 18 dancers in mind.
Don't forget the opening out of the presented 18.
Reflect on the earliest set of nine the group of 18.
Remember the starting group among the specified 18 dancers. Consider the first batch from the total 18 dancers.
Hold onto memories of the leading 9 out of the 18. Bring to mind the first subset from the given 18 performers. Think of the primary nine in the group of 18.
LIFO
Refresh your memory about the starting nine of the 18 dancers. Jog your memory concerning the first half of the specified 18. Retain thoughts of the premier nine from the ensemble of 18. Retrospect on the foremost dancers among the 18.
Focus on the first cluster of nine from the total 18.
Ponder the initial dancers in the lineup of 18.
Recollect the opening segment of nine from the presented 18. Keep in your thoughts the primary set of dancers in the group of 18.
Consider the inaugural nine performers from the indicated 18.
Remember the principal 9 dancers from the ensemble of 18.
Think of the final nine from the group 18 dancers.
Recall the concluding nine performers of the 18. Bring mind the last group among the dancers. Reflect on the 18 dancers particularly the last nine.
Keep in mind the nine dancers who concluded the group of 18. Remember the dancers who came last in the set of 18. Consider the final segment of dancers in the 18. Focus the last nine of the total 18 dancers.
Who were the concluding performers in our set of 18? Cast your thoughts the tail end of the 18 dancers. Think back to the 18 dancers especially the concluding nine. Recall the tail end of the 18-dancer lineup. Retain memories of the nine dancers wrapping up the list of 18. Mull over the dancers that concluded our 18-person roster. Bring forward the memory of the concluding acts among the 18. Remember the last stretch of dancers in the 18. Who rounded off the 18-dancer performance?
FIFO
Reminisce about the final cluster among the 18 dancers. Keep the concluding dancers in focus from the group of 18. Recapture the last performers from the total of dancers.
Table 4: task descriptions LIFO and FIFO tasks.
26

---

Published as a conference paper at ICLR 2024
Tasks
Descriptions
Recall the dancers in the room the agent frequented the most.
Think of the dancers in the room that the agent visited most often.
Bring to mind the performers in the agent's ' most-visited room. Remember the dancers from the room where the agent spent the most time.
Reflect on the dancers in the room that seemed to be the agent's 's favorite.
Consider the dancers in the chamber the agent went to repeatedly. Who were the dancers in the room the agent kept returning to? Jog your memory about the dancers the most visited room by the agent. Ponder on the performers in the room the agent seemed to prefer.
Retrospect on the dancers the space the agent was drawn to the most. Keep in mind the dancers from the room the agent visited most frequently. Recollect the dancers in the room that caught the agent's attention the most. Foeus the dancers from the room the agent seemed most interested in. Retain memories of the performers where the agent made the most stops. Mull over the dancers in the room the agent returned time and again. Reflect upon the dancers from the agent's ss most frequented chamber.
LVFO
Dwell on the performers in the room the agent seemed most attached Hold the thought of the dancers in the room with the highest visits by the agent. Contemplate dancers in the room the agent seemed to gravitate towards the most. Remember the performers in the space the agent appeared to favor.
Recall the most recent dancers from every room.
Think of the dancers who performed last in each space.
Bring to mind the final performers in every chamber.
Reflect on the concluding dancers from all rooms. Consider the last dancers to grace each stage within the rooms.
Who were the concluding performers each of the rooms? Ponder on the dancers who wrapped up the acts in every room.
Jog your memory about the closing performers from each chamber. Recollect the tail-end dancers in every room.
Focus on the final acts of dancers in all the rooms. Retain memories of the performers who closed the show in each space. Mull over the dancers who rounded off the performances in every chamber.
MVFO
Contemplate the last artists to take the stage in each of the rooms.
Think back to the closing acts from every performance space.
Zero in on the concluding dancers in all the chambers.
Dwell upon the artists who were last to perform in every room.
Cast your mind the end performers in each of the spaces.
Reflect upon the dancers who ended the sequence in each chamber.
Reminisce about the performers who took the final bow in all rooms.
Remember the artists who concluded the dances in every individual space.
Table 5: The task descriptions for LVFO and MVPO tasks.
27


Reviews:
[{"review_id": "46bvrbUyVU", "reviewer": "Reviewer_BEnq", "summary": "The paper introduces and evaluates a novel framework comprising the Spatially-Aware Transformer (SAT) and Adaptive Memory Allocator (AMA) for a range of tasks demanding spatial awareness and effective memory management. Through a series of well-structured experiments, the authors demonstrate the capabilities of their proposed methods in various environments, such as the Room Ballet environment for prediction tasks, as well as in action-conditioned world modeling and spatially-aware image generation. The comprehensiveness of the methodology is evident, as it details how to incorporate spatial information into transformer models and effectively manage memory for different types of tasks. The experiments are extensive and cover different scenarios to validate the efficacy of the proposed framework.", "questions": "See above", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "**Great Motivation:** The paper addresses a critical gap in existing models’ inability to effectively integrate spatial information, which is very crucial for tasks in various domains. The introduction and literature review provides a compelling argument for why this integration is necessary, setting a solid foundation for the rest of the paper.\n\n**Comprehensive Method Explanation:** The authors provide a thorough and clear explanation of the Spatially-Aware Transformer and Adaptive Memory Allocator. The methodology section is well-structured, detailing each component of the system, the underlying theory, and the implementation specifics, which aids in the reproducibility of the results.\n\n**Extensive Experiments:** The paper goes beyond theoretical claims and validates the proposed framework through a series of diverse and challenging experiments. These experiments not only demonstrate the strengths of the SAT-AMA combination but also highlight its versatility across different tasks and scenarios. The image generation experiments are especially interesting and seem not to exist in its ancestor work of Towards mental time travel: a hierarchical memory for reinforcement learning agents.", "weaknesses": "**Usability in Complex Embodied AI Scenarios:** The paper, while comprehensive, could benefit from a deeper discussion on the applicability and scalability of the proposed methods in more complex embodied AI scenarios. Given the rising interest in virtual homes and ThreeDWorld with many rooms for task operation, readers would appreciate some discussion or insights into how the SAT-AMA framework can be adapted or scaled to meet the challenges presented by these intricate environments. A similar discussion could be like, \"Is the method scalable to larger environments?\" and \"How can the SAT-AMA framework be adapted for more intricate task operations?\" to provide a complete picture to the readers.\n\n**More discussion on AMA:** I appreciate the introduction and application of the Adaptive Memory Allocator (AMA) within the Spatial Awareness Transformer framework, as it presents a novel and promising approach to dynamically allocate memory based on task requirements. However, I find that there could be a more detailed analysis and discussion of AMA’s strategy selection across different experimental settings and tasks. Understanding how AMA decides on specific spatial strategies could uncover valuable priors for selecting appropriate strategies based on the nature of the task, which would immensely benefit future work exploring new methods in this domain. It would be beneficial for the readers if the authors could provide insights into the distribution of strategies chosen by AMA, its dependency on the nature of tasks, and the correlation between strategy choices and task performance. Such analysis would not only deepen our understanding of AMA’s workings but also offer practical guidance for researchers aiming to employ similar adaptive mechanisms in their work. A discussion on these aspects would significantly enhance the completeness and depth of the paper, providing valuable context and potentially guiding future innovations in the field.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HCcdCy9cfh", "reviewer": "Reviewer_XL1m", "summary": "Inspired by the significance of spatial context in the formation and retrieval of episodic memory, the paper proposes to add a spatial embedding into transformers and organize the episodic memory in a place-centric way to obtain a spatially aware Transformer model. The paper also proposes the Adaptive Memory Allocator, a memory management method based on reinforcement learning that aims to optimize the efficiency of memory utilization. The experiments on several environments demonstrate the advantages of the proposed model.", "questions": "See concerns in the weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- It's exciting and of significance to have a transformer capable of utilizing explicit spatial information and can act as better episodic memory.\n\n- Designing an adaptive memory allocator is useful\n\n- Extensive experiments on various environments and tasks.\n\n- The idea of incorporating spatial information into episodic memory is novel and interesting", "weaknesses": "- It seems the \"spatial-aware transformers\" are achieved solely by adding a spatial embedding. And its specific design or implementation is not clearly mentioned. Only a sinusoidal positional embedding is mentioned in Exp-1, do other experiments also use this? How is this positional embedding enough to represent spatial relations, since space is not unidirectional as time? \n\n- The idea of an ADAPTIVE MEMORY ALLOCATOR is very interesting, but the implementation is quite trivial to me. What's the difficulty of learning this policy through Q-learning?\n\n- 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.\n\n- Much more details of the experiments are needed. What are the exact input and output of these environments? How is the input represented and fed into the transformers? How is the training implemented for different baselines?\n\n- The experiment environments are too toy setting for \"embodied agents\", it would strengthen the paper to add some real embodied environments, such as the Habitat[1] and TDW[2].\n\n[1] Habitat: A Platform for Embodied AI Research\n\n[2] ThreeDWorld: A Platform for Interactive Multi-Modal Physical Simulation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KAfk4GBkKo", "reviewer": "Reviewer_QvRD", "summary": "The paper investigates the spatially-aware transformers (SAT) that incorporate agent's spatial and temporal experiences important to solve embodied problems. Specifically, it studies the influence of spatial and decisional information, and explores multiple memory management strategies (e.g., Place-centric vs Time-centric, First-In-First-Out vs Last-In-First-Out, and Adaptive Memory Allocator) on various downstream tasks, showing insights of different memory utilization and management approaches in addressing many machine learning problems (e.g., supervised prediction, image generation), and justifying the proposed SAT-AMA.", "questions": "I hope the authors can address some of my concerns in Weaknesses. I don't have any other questions here.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- This paper studies the important problem of utilizing and managing spatial and temporal memory experienced by embodied agents that are essential to address various downstream tasks. Particularly, it considers the practical issue of memory constraint and perform experiments based on the popular transformer architectures. The research presented in this paper is very well motivated.\n\n- The paper is very technical solid. Almost all arguments are well-supported/justified by the highly-relevant and classic references and thorough experiments (in Appendix). It carefully compares place- and time-centric store and read, and progressively studies FIFO to the proposed Adaptive Memory Allocator based on SAT. I believe the extensive settings and results presented in this paper have the potential to inspire many future works.\n\n- Overall, this paper was a very enjoyable read to me. All details have been clearly presented (especially with the Appendix and all nice visualizations); The paper is nicely-structured, concise but contains massive valuable information. I believe many arguments and thoughts presented in this paper will be very constructive to relevant future research.", "weaknesses": "- The title of this paper is very misleading\n    - 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. \n\n- This paper overclaims several contributions.\n    - 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.\n    - 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. \n    - 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.\n\n- Missing experiments.\n    - I think this paper lacks investigation on the more recent embodied agents and their memory management approaches. \n    - 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? \n    - Following my previous point, I think this study emphasizes the scenario of memory constraint, but the experiments only considers small memory capacity, negelating 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.\n    - 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).\n\n[1] Think Global, Act Local: Dual-scale Graph Transformer for Vision-and-Language Navigation. Chen et al., CVPR2022.\n\n[2] Vision-and-Language Navigation: Interpreting Visually-Grounded Navigation Instructions in Real Environments. Anderson et al., CVPR2018.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PkPYVhbNEz", "reviewer": "Reviewer_44cr", "summary": "The paper proposes Spatially-Aware Transformers (SAT) to keep spatial information in place-centric episodic memory.\nSAT maintains multiple spatial memories for respective places while ensuring the \"total\" size of memory (i.e., $L / K$).\nFor memory management, the author proposes the Adaptive Memory Allocator (AMA) that adaptively finds the best memory management policy by learning $\\pi(\\sigma | \\tau)$ using Q-learning given a task description, $\\tau$.\nThe proposed AMA outperforms the baselines (i.e., model w/o AMA) by noticeable margins regarding effectiveness and efficiency in various downstream tasks.", "questions": "See weaknesses above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is generally written well and easy to follow.\n- Extending Transformer-based memory architecture to spatial domains is well-motivated and sounds sensible.\n- Learning to choose the best memory management strategy from multiple candidates looks reasonable.\n- Experiments on various downstream tasks supports the generality of the proposed approach.\n- The proposed approach achieves strong performance gain with large margins.", "weaknesses": "- The novelty of AMA seems a bit weak, as it is basically policy learning that chooses the best action (here, strategy) that maximizes rewards. What are some core differences from conventional policy learning, especially related to memory management for spatial information?\n- 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)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)