Research map / Architectures and efficiency / Attention and transformers
Spatially-Aware Transformers for Embodied Agents
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Spatially-Aware Transformers (SAT) that add explicit spatial embeddings to episodic transformer memory, with place-centric hierarchical memory and an RL-based Adaptive Memory Allocator (AMA) that selects among memory write strategies, evaluated on Room Ballet, FFHQ navigation generation, MiniWorld, and a small RL task.
Problem
Transformer-based episodic memory only uses temporal order (FIFO with positional embeddings), ignoring spatial context that is important for place-centric recall and efficient memory management in embodied agents.
Approach
Add explicit spatial (place) embeddings to memory frames, organize memory into per-place chunked hierarchies with place-centric top-k reading, and learn via one-step Q-learning to choose among memory write/removal strategies (AMA) conditioned on the task description.
Key results
- On the Next Ballet spatial reasoning task, T-FIFO fails completely while SAT-FIFO solves it, showing explicit spatial info is needed.
- SAT-AMA trained with language task descriptions achieves over 95% accuracy on unseen description variants in Exp-4.
- SAT-PM-PH consistently outperforms SAT-FIFO-TH and SAT-FIFO in the Short Stay task, with the gap widening as room count increases.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add explicit spatial information to the transformer | SAT-FIFO · Room Ballet Next Ballet | accuracy | Helped |
| add action information without explicit spatial information (ablation) | T-FIFO · Room Ballet Next Ballet | accuracy | Helped |
| use place-centric chunks instead of time-centric chunks (ablation) | SAT-PM-PH · Room Ballet Short Stay | accuracy | Helped |
| use place-centric hierarchical read instead of reading the full memory (ablation) | SAT-PM-PH · Room Ballet Short Stay | inference speed | Helped |
| partition memory by place rather than using flat FIFO memory (ablation) | SAT-PM · Room Ballet Ballet-ABA | accuracy | Helped |
| select memory allocation strategy adaptively for the task | SAT-AMA · Room Ballet Ballet-Strategy tasks | accuracy | Helped |
| select memory allocation strategy from natural-language task descriptions | SAT-AMA-Lang · Room Ballet multi-task descriptions | accuracy: over 95% | Helped |
| use SAT-FIFO instead of SAT-AMA's adaptive strategy selection (ablation) | MiniWorld image generation · MiniWorld | generation performance | Hurt |
| use place-centric hierarchical reading instead of time-centric hierarchical reading (ablation) | SAT-PM-PH · SAT Read inference | inference time (ms): 7.51 \pm 1.58 -> 7.49 \pm 1.29 | No clear effect |
| use FIFO rather than LIFO, MVFO, or LVFO | AMA on Ballet-FIFO · Ballet-FIFO | accuracy: 0.146, 0.325, 0.714 -> 0.999 | Helped |
Surprising findings
- Transformers with only temporal order and actions (T-FIFO and T-FIFO+A) fail or perform poorly on the spatial reasoning task even with action information.
- The paper reports that training AMA through learning over all write decisions at each step is notoriously difficult, motivating the restricted strategy-selection design.
Limitations
- Relies on explicit spatial annotations being available.
- AMA is limited to a predefined set of memory strategies rather than learning strategies from scratch.
- Evaluated only on spatial reasoning tasks in simulated environments.
Experimental setup
Datasets: FFHQ Room Ballet (simulated) MiniWorld (simulated)
Benchmarks: Room Ballet Next Ballet Room Ballet Short Stay Ballet strategy tasks (FIFO/LIFO/MVFO/LVFO) FFHQ navigation generation MiniWorld generation Visual Match on MiniWorld two-room RL task
Metrics: accuracy generation quality inference time (ms) box collection success
Baselines: T-FIFO T-FIFO+A SAT-FIFO-TH SAT-FIFO DNC PPO
Builds on: Transformer-XL Hierarchical Chunk Attention Memory (HCAM) Differentiable Neural Computer Proximal Policy Optimization
Closest papers by meaning
- Memo: Training Memory-Efficient Embodied Agents with Reinforcement Learning (NeurIPS 2025 · similarity 0.89)
- Episodic Memories Generation and Evaluation Benchmark for Large Language Models (ICLR 2025 · similarity 0.88)
- 3D-Mem: 3D Scene Memory for Embodied Exploration and Reasoning (CVPR 2025 · similarity 0.86)
- Video World Models with Long-term Spatial Memory (NeurIPS 2025 · similarity 0.86)
- Embodied Agents Meet Personalization: Investigating Challenges and Solutions Through the Lens of Memory Utilization (ICLR 2026 · similarity 0.86)
- Planning from Imagination: Episodic Simulation and Episodic Memory for Vision-and-Language Navigation (AAAI 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.