Source-linked AI summary
Neural Map: Structured Memory for Deep Reinforcement Learning
Emilio Parisotto, Ruslan Salakhutdinov
TL;DR
DRL agents need memory for planning in partially observable environments, while prior approaches are limited by short or fixed memory. The paper introduces the spatially structured Neural Map with adaptable writes and evaluates it in 2D and 3D mazes, where it surpasses baseline memories and generalizes to unseen environments.
Problem
Prior DRL memories rely on short temporal context or fixed recent-state pools, limiting retention of information over long time lags.
Method
The Neural Map uses an adaptable sparse write operation in a spatially structured 2D memory whose addresses correspond to the agent’s position.
Results
The Neural Map surpasses LSTM and memory-network baselines across challenging 2D mazes and also outperforms other methods in a Doom-based 3D maze, including unseen maps.
Takeaways & Limitations
Spatially structured memory with learned writes supports long-term maze behavior and generalization beyond training environments.
Takeaways & Limitations
The Neural Map requires an oracle for the agent’s current position, while its ego-centric alternative requires extracting the agent’s velocity.
Abstract
from arXiv · showhide
A critical component to enabling intelligent reasoning in partially observable environments is memory. Despite this importance, Deep Reinforcement Learning (DRL) agents have so far used relatively simple memory architectures, with the main methods to overcome partial observability being either a temporal convolution over the past k frames or an LSTM layer. More recent work (Oh et al., 2016) has went beyond these architectures by using memory networks which can allow more sophisticated addressing schemes over the past k frames. But even these architectures are unsatisfactory due to the reason that they are limited to only remembering information from the last k frames. In this paper, we develop a memory system with an adaptable write operator that is customized to the sorts of 3D environments that DRL agents typically interact with. This architecture, called the Neural Map, uses a spatially structured 2D memory image to learn to store arbitrary information about the environment over long time lags. We demonstrate empirically that the Neural Map surpasses previous DRL memories on a set of challenging 2D and 3D maze environments and show that it is capable of generalizing to environments that were not seen during training.
1 INTRODUCTION
Memory enables planning and reasoning beyond immediate percepts, but existing DRL memories either use short temporal context or fixed memory contents. The Neural Map introduces spatially structured, adaptable memory and is evaluated in 2D and 3D mazes.
- Motivation: Long-term planning tasks require agents to retain information beyond immediate observations.Examples include gathering objects and searching for a key in a role-playing game.
- Existing memories: Memory Networks learn how to read a fixed pool of recent states rather than which contents to store.Their memory typically contains the past M states observed in the environment.
- Existing memories: Fixed-memory approaches can store redundant information and require experts to set a memory size beyond the task’s time horizon.The memory size M must be chosen in advance for the task under consideration.
- Write-based memories: Write-based memories can learn to retain salient information for unbounded time steps while ignoring useless information.This avoids requiring a priori knowledge of what should be stored.
2 BACKGROUND
The paper frames reinforcement learning with Markov decision processes, policies, value functions, and actor-critic optimization. Its implementation uses a synchronous modification of A3C across concurrent environments.
- MDP formulation: An MDP is defined by states, actions, transition probabilities, a discount factor, and rewards.The policy maps states to action distributions.
- MDP formulation: The optimal value function is the maximum policy value achievable from each state.An optimal policy attains this value at every state.
- Policy optimization: REINFORCE updates a policy using the log-probability gradient weighted by future cumulative reward.The update can use a value-function baseline to reduce variance.
- Policy optimization: The resulting actor-critic formulation combines policy-gradient updates with a value-function baseline.The baseline-augmented update uses G_t − V^π(s_t).
- A3C implementation: The paper uses a modified A3C framework in which deep networks parameterize the policy and value function.The policy is represented by a deep neural network with parameters θ.
- A3C implementation: Unlike standard asynchronous A3C, the implementation updates policies synchronously at every time step across concurrent environments.Network parameters therefore never require synchronization.
3 NEURAL MAP
The Neural Map is an internal memory that stores and reads features in a spatially organized map. Its selective write operation updates only the map region corresponding to the agent’s current location.
- Memory structure: The Neural Map is readable and writable internal storage whose writes are restricted to the agent’s current map location.The paper uses a 2D map for simplicity and notes that higher-dimensional extensions are possible.
- Memory structure: The map can be extended from a 2D representation to 3D or higher-dimensional structures.One example is a 4D map containing a 3D sub-map for each cardinal direction.
- Map representation: The neural map M is a C × H × W feature block, with C feature channels and H and W denoting its vertical and horizontal extents.Agent coordinates are normalized into positions within the map.
- Map operations: At time t, the Neural Map combines the current state embedding, map, and agent position to produce its output.The output is subsequently processed by a deep network to produce policy outputs.
3.1 GLOBAL READ OPERATION
The Neural Map reads global information through convolution and performs context-based associative addressing over map locations. Soft attention weights produce a weighted memory summary for downstream policy computation.
- Global read: The global read applies a deep convolutional network to the map and produces a feature vector summarizing the entire map.This vector is denoted r_t.
- Context read: The context operation combines the current state embedding with the global read to form a query for map addressing.The query is compared with features stored at map positions.
- Context read: Scores for all map positions are normalized into a probability distribution over locations.This distribution is the paper’s soft-attention mechanism.
- Context read: The attention distribution computes a weighted average of features across the map.The resulting context vector summarizes the retrieved map information.
- Context read: Associative addressing lets the agent query incomplete information and retrieve the stored memory that most closely matches it.The paper gives querying whether a viewed landmark has been seen before as an example.
3.3 LOCAL WRITE OPERATION
The local write operation uses the agent’s current position and local map feature together with state and read information to produce a new feature vector. A deep neural network generates the candidate written at that position.
- At position (x_t, y_t), the write operation receives the state embedding, global read output, context read vector, and current map feature.
- A deep neural network f transforms these inputs into a new C-dimensional vector w^(x_t,y_t)_(t+1).
- The generated vector serves as the local write candidate for the agent’s current position.
3.4 MAP UPDATE OPERATION
The map update preserves the existing neural map except at the agent’s current position, where it applies the current write candidate vector.
- The next map M_(t+1) equals M_t everywhere except the agent’s current position.
- At the current position, the update replaces the stored feature with the current write candidate vector.
- Thus, each update changes only the location currently occupied by the agent.
3.5 OPERATION VARIANTS
The paper describes variants that reduce map-processing cost, separate addressing from content, and replace hard local rewrites with gated updates. These modifications tailor map operations to computational constraints and information retention.
- Operation variants: The section introduces these changes as variants of the standard neural-map operations.
- Spatial subsampling: A spatial subset of the neural map can be processed instead of the entire map, reducing the cost of handling large high-resolution maps.A Spatial Transformer Network can attentively subsample particular locations and scales.
- Key-value context read: Splitting each map feature into key and value parts strengthens context addressing by matching keys to the query and averaging the values.This separates the addressing feature space from the content feature space.
- GRU-based local write: A GRU-based write operation replaces hard rewriting with reset and update gates that modulate how much the new vector differs from the stored feature.The formulation uses recurrent update equations and defines reset, candidate-activation, and update components.
4 EXPERIMENTS
The experiments evaluate Neural Map agents against LSTM and memory-network baselines in partially observable 2D and 3D maze environments. Neural Maps solve more mazes, handle larger mazes better, and the GRU-based update improves learning speed and stability.
- 2D Goal-Search Environment: Neural Map architectures solve the most Goal-Search mazes across both training and held-out test distributions.The GRU-based Neural Map solves almost all training and test mazes.
- 2D Goal-Search Environment: Memory-network and LSTM agents struggle on larger mazes, whereas both standard and GRU-based Neural Maps solve them at much higher rates.The evaluation separates held-out mazes into small and large size ranges to assess longer information-retention requirements.
- 2D Goal-Search Environment: The GRU-based Neural Map learns faster, reaches a higher final score, and is typically more stable than the standard-update Neural Map.Both Neural Map variants initially learn more slowly than the LSTM and MemNN, but achieve the highest final reward.
- 2D Goal-Search Environment: Context retrieval revisits the map location of the previously observed indicator when the agent encounters goal colors.The retrieved location becomes especially strong when the agent first observes the correct teal goal, showing use of a stored landmark.
- 3D Doom Environment: Standalone Neural Map results are omitted in Doom because coarse spatial granularity caused repeated actions and oscillation near walls.An added pre-output LSTM tracks recent actions and helps the agent turn consistently.
- 3D Doom Environment: In the 3D Doom maze, the LSTM plus GRU-based Neural Map surpasses other tested methods on both the training map and six unseen maps.Episodes require finding the tower matching a colored indicator within 420 steps, with randomly sampled goal locations.
5 EXTENSION: EGO-CENTRIC NEURAL MAP
The ego-centric Neural Map avoids requiring absolute agent coordinates by shifting the map according to measured velocity, keeping the agent centered while the world moves around it.
- Absolute-position Neural Maps require an oracle for the agent’s current (x, y) position, which is difficult and remains unsolved.
- Relative positioning counter-transforms the map by the inverse velocity so the agent remains stationary at its center.The method assumes a function ξ discretizes measured velocities into valid Neural Map velocities.
- The approach handles reverse translation directly and could also handle rotations when angular velocity is measurable.
- The ego-centric map uses the center position (p_w, p_h) as the sole write location during map updates.
6 RELATED WORK
Prior DRL memory approaches include fixed-memory networks and recurrent architectures modeled on computer memory, while other systems use engineered mapping pipelines or task-specific spatial memories.
- Memory Networks store a fixed set of recent states and learn how to read them rather than what information to store.MemNN encodes the last M states as key-value pairs and retrieves them with soft attention.
- Computer-inspired recurrent architectures explicitly separate neural computation from external memory, including systems with multiple read and write mechanisms.The DNC differs from the Neural Map by using an unconstrained write operator rather than a fixed write location.
- A prior 3D mapping approach augments DRL inputs through hard-coded modules such as SLAM and image segmentation, whereas the Neural Map is trained end-to-end without weak supervision.
- Concurrent spatial-memory work applied a 2D map structure to robot navigation in quantized office-building grid worlds, while the Neural Map is presented as a general memory architecture.
7 CONCLUSION
The paper presents a spatially structured Neural Map with sparse position-corresponding writes, evaluates it on long-horizon 2D mazes, and extends it to Doom through a hybrid Neural Map–LSTM model.
- The Neural Map organizes memory as a 2D spatial map and writes sparsely at addresses corresponding to the agent’s current position.
- In a challenging 2D maze task requiring long-term information storage, the architecture surpassed baseline memories from previous work.
- The GRU-based update equation was crucial for improving both learning speed and training stability.
- A hybrid Neural Map + LSTM model solved most Doom scenarios and surpassed both LSTM and MemNN baseline agents.