Technical Summary: A Machine with Short-Term, Episodic, and Semantic Memory Systems

Problem Statement

Current machine learning agents, including search engines and virtual assistants, excel at retrieving factual, commonsense knowledge (semantic memory) but lack the ability to recall specific past events or personal experiences (episodic memory). Furthermore, standard reinforcement learning (RL) agents often struggle with partial observability and catastrophic forgetting. The authors posit that mimicking the explicit human memory architecture—specifically the distinction between semantic (general world knowledge) and episodic (personal, time-bound experiences) memory, mediated by a short-term buffer—could enable agents to better manage limited memory capacity and answer diverse questions more effectively.

Methodology

The Environment: "The Room"

To evaluate this hypothesis, the authors designed a discrete-event simulation environment compatible with OpenAI Gym.

• Setup: The environment contains $N_{humans}$ (64), $N_{objects}$ (16), and $N_{locations}$ (28). Humans follow fixed routines, placing objects in locations that are either "commonsense" (e.g., a bowl on a cupboard) or non-commonsense with a probability $p_{commonsense} = 0.5$.
• Observation: The agent partially observes the environment, viewing one human at a time. Observations are recorded as quadruples $(h, r, t, timestamp)$, representing a head entity, relation, tail entity, and timestamp.
• Task: At each step, the agent is asked a question in the form of a triple $(h, r, ?)$. Correct answers yield a reward of +1; incorrect answers yield 0. The episode lasts 128 steps.

Memory Architecture

The agent is modeled with three distinct memory systems, each represented as a Knowledge Graph (KG):

1. Short-Term Memory ($M_o$): A bounded buffer that initially stores all observations as quadruples $(h, r, t, timestamp)$.
2. Episodic Memory ($M_e$): A bounded long-term memory storing specific, time-bound events as quadruples $(h, r, t, timestamp)$.
3. Semantic Memory ($M_s$): A bounded long-term memory storing general facts as quadruples $(h, r, t, strength)$. Here, $h$ is generalized (e.g., "laptop" rather than "Bob's laptop"), and $strength$ indicates the frequency of storage.

Memory Management Policy:
When the short-term memory is full, the agent must decide the fate of the oldest observation. It has three actions:

1. Forget: Discard the observation entirely.
2. Store in Episodic: Move the observation to $M_e$.
3. Store in Semantic: Move the observation to $M_s$ (generalizing the entity and incrementing strength).

The retrieval policy is deterministic and rule-based (Algorithm 1):

• For specific questions, the agent first checks $M_e$ for the most recent relevant event.
• If no episodic match exists, it checks $M_s$ for the strongest relevant general fact.
• If neither yields an answer, the agent fails the question.

Learning Algorithm

The memory management policy (deciding where to store or whether to forget) is learned using Deep Q-Learning (DQN).

• State Representation: The state $s$ consists of the current contents of $M_o$, $M_e$, and $M_s$.
• Input Processing: To feed KGs into the neural network, the authors convert memories into Knowledge Graph Embeddings (KGE). Quadruples are sorted (by timestamp for episodic/short-term, by strength for semantic) and embedded. For episodic/short-term memories involving humans, the human name and object are concatenated before embedding to capture individual characteristics.
• Network Architecture: The Q-network utilizes LSTM layers to process the sequence of memory embeddings, followed by Multi-Layer Perceptrons (MLPs) to output Q-values for the three actions.
• Training Variants: Two agent types were trained:
  1. Semantic Scratch: Starts with an empty semantic memory.
  2. Semantic Pretrained: Starts with general world knowledge (from ConceptNet) pre-filled in $M_s$.

Key Contributions

1. Explicit Memory Modeling: The paper explicitly models an agent with short-term, episodic, and semantic memory systems, each structured as a knowledge graph, inspired by cognitive science theories (Tulving 1983, 1985).
2. Novel Environment: The release of "The Room," a custom RL environment designed to test an agent's ability to encode, store, and retrieve memories to answer questions under partial observability.
3. Learned Memory Management: Demonstration that a DQN-based agent can successfully learn a policy for managing memory (deciding between forgetting, episodic storage, or semantic storage) to maximize returns.

Results

Experiments were conducted with varying memory capacities (total capacity $|M_e| + |M_s|$ ranging from 2 to 64).

• Performance vs. Baselines: Agents with both episodic and semantic memory systems significantly outperformed three baselines: (1) Episodic-only, (2) Semantic-only, and (3) Random action selection. The best-performing agent (RL with pre-trained semantic memory) achieved an average total reward of 116.5 out of 128, compared to 109.1 for the scratch version and lower scores for baselines.
• Impact of Pre-training: The "Semantic Pretrained" agent learned faster and converged to a higher optimum than the "Semantic Scratch" agent. The pretrained agent learned to preserve its initial general knowledge and only store specific, non-generalizable events in episodic memory. Conversely, the scratch agent attempted to learn general knowledge from scratch, occasionally making errors (e.g., misidentifying a train's location) which lowered its final reward.
• Capacity Scaling:
  • Episodic-only agents showed linear improvement with capacity, eventually answering all 128 questions when capacity reached 64.
  • Semantic-only agents performed well at low capacities but plateaued (around 92 rewards) once capacity exceeded 32, as the general knowledge triples were sufficient for the task.
  • Combined agents leveraged both systems to achieve the highest overall performance.

Significance and Claims

The paper claims that equipping machines with human-like, explicit memory systems allows them to outperform agents without such structures in environments requiring both general knowledge and specific event recall. The authors emphasize that their approach successfully learns a memory management policy via reinforcement learning, effectively deciding what to forget and what to retain in different memory systems.

The work bridges cognitive architecture (which models human mind theory) and modern machine learning (which focuses on solving partial observability). The authors note that while their current environment is a discrete simulation, the framework is designed to be extensible to more complex, multimodal scenarios and natural language queries in future work. They conclude that their agent demonstrates the viability of using knowledge graphs within RL agents to manage memory constraints and improve decision-making.