# Spatially-Aware Transformer for Embodied Agents

arXiv: https://arxiv.org/abs/2402.15160 · HTML: https://mlsanity.com/paper/2402.15160/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 25 of its 45 references that are in the map.
- First posted on arXiv 23 Feb 2024.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Habitat: A Platform for Embodied AI Research](https://mlsanity.com/paper/1904.01201/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [MineRL: A Large-Scale Dataset of Minecraft Demonstrations](https://mlsanity.com/paper/1907.13440/) (2019)
- [TransDreamer: Reinforcement Learning with Transformer World Models](https://mlsanity.com/paper/2202.09481/) (2022)
- [Shaping Belief States with Generative Environment Models for RL](https://mlsanity.com/paper/1906.09237/) (2019)
- [Relating transformers to models and neural representations of the hippocampal formation](https://mlsanity.com/paper/2112.04035/) (2021)
- [Optimizing Agent Behavior over Long Time Scales by Transporting Value](https://mlsanity.com/paper/1810.06721/) (2018)
- [Sequential Neural Processes](https://mlsanity.com/paper/1906.10264/) (2019)
- [Towards mental time travel: a hierarchical memory for reinforcement learning agents](https://mlsanity.com/paper/2105.14039/) (2021)
- [Efficient Transformers in Reinforcement Learning using Actor-Learner Distillation](https://mlsanity.com/paper/2104.01655/) (2021)
- [Robustifying Sequential Neural Processes](https://mlsanity.com/paper/2006.15987/) (2020)
- [Generative Temporal Models with Spatial Memory for Partially Observed Environments](https://mlsanity.com/paper/1804.09401/) (2018)
- [Learning models for visual 3D localization with implicit mapping](https://mlsanity.com/paper/1807.03149/) (2018)

## Built on by (most cited first, 5 of 5)

- [Visual Embodied Brain: Let Multimodal Large Language Models See, Think, and Control in Spaces](https://mlsanity.com/paper/2506.00123/) (2025)
- [Neural Brain: A Neuroscience-inspired Framework for Embodied Agents](https://mlsanity.com/paper/2505.07634/) (2025)
- [Mind Meets Space: Rethinking Agentic Spatial Intelligence from a Neuroscience-inspired Perspective](https://mlsanity.com/paper/2509.09154/) (2025)
- [Aligning Perception, Reasoning, Modeling and Interaction: A Survey on Physical AI](https://mlsanity.com/paper/2510.04978/) (2025)
- [Spatially-informed transformers: Injecting geostatistical covariance biases into self-attention for spatio-temporal forecasting](https://mlsanity.com/paper/2512.17696/) (2025)

## Abstract (arXiv)

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 https://github.com/junmokane/spatially-aware-transformer.