AriGraph: Learning Knowledge Graph World Models with Episodic Memory for LLM Agents

1 Introduction *

Impressive language generation capabilities of large language models (LLMs) has sparked substantial interest in their application as core components for creating autonomous agents capable of interacting with dynamic environments and executing complex tasks. Over the past year, the research community has explored general architectures and core modules for such LLM agents [Wang et al., 2024; Sumers et al., 2024; Cheng et al., 2024]. A crucial property of a general cognitive agent is its ability to accumulate and use knowledge. A

* An extended version of the paper, including appendices, is available at arXiv:2407.04363.

long-term memory allows an agent to store and recall past experiences and knowledge, enabling it to learn from previous encounters and make informed decisions. However, the question of the best way to equip an agent with these capabilities remains open. Despite the constraints inherent in transformer architectures, contemporary methods enable LLMs to manage contexts encompassing millions of tokens [Ding et al., 2024b]. However, this approach proves inefficient for agents required to maintain continuous interaction with their environment. Such agents must hold an entire historical context in memory to perform actions, which is not only costly but also limited in handling complex logic hidden in vast amounts of information. Research into alternative frameworks like Recurrent Memory Transformer [Bulatov et al., 2022; Bulatov et al., 2024] and MAMBA [Gu and Dao, 2023] seeks to provide long-term memory solutions, though these models are still in their infancy.

Currently, the most popular solution for incorporating memory to LLM agents is the Retrieval-Augmented Generation (RAG) approach. RAG in a form of vector retrieval leverages an external database to enhance the model's prompt with relevant information. This technique is commonly used in memory architectures for LLM agents, often to recall specific observations or learned skills. However, it suffers from unstructured nature, greatly reducing the ability to retrieve related information, which may be scattered throughout the agent's memory. These limitations can be overcome by using knowledge graphs as database. This approach has also experienced a resurgence with the advent of LLMs [Pan et al., 2024]. However, for a robust memory architecture, integrating both structured and unstructured data is essential. In cognitive science, this integration parallels the concepts of semantic and episodic memories. Semantic memory encompasses factual knowledge about the world, whereas episodic memory pertains to personal experiences, which often contain richer and more detailed information. Though traditionally considered separate due to their distinct neurological representations, recent studies suggest these memory types are interconnected [Wong Gonzalez, 2018]. Semantic knowledge is built upon the foundation of episodic memory and subsequently provides a structured base for associative memory. This allows for the integration of various memory aspects, including episodic memories themselves.

In our research, we have developed a memory architecture called Ariadne's Graph (AriGraph), that integrates semantic and episodic memories within a memory graph framework. A knowledge graph represents a network of interconnected semantic knowledge, while episodic memories are depicted as episodic edges that can connect multiple relations within the graph. As an agent interacts with environment, it learns joint semantic and episodic world model by updating and extending knowledge graph based memory. This architecture not only serves as a foundational memory framework but also aids in environmental modeling, improving spatial orientation and exploration capabilities. For the general framework of our LLM agent called Ariadne, we employed pipeline of memory retrieval, planing and decision making. For evaluation of proposed methods we set up experiments to study two research questions.

• RQ1. Can LLM based agents learn useful structured world model from scratch via interaction with an environment?

RQ2. Does structured knowledge representation improve retrieval of relevant facts from memory and enable effective exploration?

We evaluated our agent in complex interactive tasks in Textworld and NetHack environments [Cˆot´e et al., 2018; K¨uttler et al., 2020]. Experimental results demonstrate that our agent Ariadne can effectively learn through interactions with environment and significantly outperforms other memory approaches for LLMs such as full history, summarization, RAG, Simulacra [Park et al., 2023] and Reflexion [Shinn et al., 2023]. We also show that our method outperforms existing reinforcement learning (RL) baselines. We also evaluated our approach on the classical roguelike game NetHack, where our agent with local observations achieved scores comparable to an agent with ground-truth knowledge. Although AriGraph was originally designed for an agent interacting with the environment, it also demonstrates competitive performance on multi-hop question answering tasks.

2 AriGraphWorldModel

Memory graph structure. AriGraph world model G = (Vs,Es, Ve,Ee) consists of semantic (Vs,Es) and episodic memory (Ve,Ee) vertices and edges (see Figure 2). At each step t agent receives observation ot and sends action at back to the environment. The environment also returns rewards rt that are not visible to the LLM agent but are used to evaluate its performance. The agent continuously learns world model G by extracting semantic triplets (object1, relation, object2) from textual observations ot.

• Vs is a set of semantic vertices. Semantic vertices correspond to objects extracted from triplets.
• Es is a set of semantic edges. Semantic edge is a tuple (v, rel, u), where u, v are semantic vertices and rel is a relationship between them. Semantic edges essentially represent triplets integrated in the semantic memory.
• Ve is a set of episodic vertices. Each episodic vertex corresponds to an observation received from the environment at the respective step vte = ot.
• Ee is a set of episodic edges. Each episodic edge ete = (vte , Ets ) connects all semantic triplets Ets extracted from ot with each other and corresponding episodic ver-

Algorithm 1: Memory Graph Search

Input: set of queries Q, Vs, Es, Ve, Ee,

number of episodic vertices k, semantic search depth d and width w

Result: retrieved episodic vertices V Q e , retrieved semantic triplets EQ s

foreach q in Q do

end

tex vte . In other words episodic edges represent temporal relationship "happened at the same time". 1

Constructing AriGraph. Interaction with the environment can provide the agent with an information about the world to create new or update previously acquired knowledge. Given new observation ot, LLM agent extracts new triplets as semantic vertices V t s and edges Ets . To find already existing knowledge about the objects mentioned in ot a set of all semantic edges Erel s incident to vertices V t s is filtered out. Then outdated edges in Erel s are detected by comparing them with Ets and removed from the graph. After clearing outdated knowledge we expand semantic memory with V t s and Ets . Episodic memory is updated by simply adding new episodic vertex vte containing ot and new episodic edge that connect all edges in Ets with vte . Episodic nodes store agent's past history and episodic edges connects all knowledge received at the same step. See Appendix E for prompts used to extract new triplets and detect outdated knowledge.

Retrieval from AriGraph. For successful decisionmaking in a partially observable environment, the agent needs to be able to retrieve relevant knowledge. Retrieval from the AriGraph memory consists of two procedures: (1) a semantic search returns the most relevant triplets (semantic edges) and (2) an episodic search that, given extracted triplets, returns the most relevant episodic vertices Ve. The pseudo-code for the search is presented in the Algorithm 1.

Semantic search relies on semantic similarity and semantic graph structure to recall the most relevant triplets. Given a query, the retriever (pre-trained Contriever model [Izacard et al., 2022]) selects the most relevant semantic triplets. Then, the set of vertices incident to the found triplets is used to recursively retrieve new edges from the graph. Depth and breadth of the search can be controlled by respective hyperparameters d and w. For details see Appendix A.

Episodic search starts with the results of the semantic search as an input. Episodic edges link the input triplets with past episodic vertices representing observations. The number of input triplets associated with a particular episodic vertex is used to calculate their relevance:

1 Strictly speaking, episodic edges cannot be called edges or even hyperedges, because they connect vertices with multiple graph edges, but for simplicity we call them edges or episodic edges.

where ni is a number of input triplets incident to episodic edge ei , Ni is a total number triplets (semantic edges) incident to ei and log (max(Ni, 1)) is a weighting factor to prevent high scores for low information observations. k most relevant episodic vertices (containing respective observations) are returned as a result of the episodic search.

3 Ariadne Cognitive Architecture

To test utility of AriGraph world modelling method we propose an agentic architecture called Ariadne. Ariadne agent interacts with an unknown environment to accomplish a goal set by a user. Throughout this process, at each time step, the agent learns a world model, plans and executes actions. Ariadne has long-term memory stored as AriGraph and working memory containing information for current planning and decision making.

Given an observation the agent updates world model and retrieves semantic and episodic knowledge from AriGraph to working memory. Working memory is also populated with a final goal description, current observation, history of recent observation and actions. At the planning stage, Ariadne agent uses content of working memory to create new or update existing plan as a series of task-relevant sub-goals, each accompanied by a concise description. The planning module also evaluates the outcomes of actions based on feedback from the environment after each action at step t - 1, adjusting the plan accordingly.

The revised plan is added to the working memory which is accessed by the decision-making module, tasked with selecting the most suitable action aligned with the current plan's objectives. This module adheres to the ReAct [Yao et al., 2023] framework, requiring the agent to articulate the rationale behind an action before execution. Separation of planning from decision-making enables LLMs to focus on distinct cognitive processes. In text-based environments an agent selects an action from the list of valid actions. Our agent can also use graph specific function for navigation utilizing its memory module. It extends its action space with "go to location" type commands and infers an optimal route to a target location using spatial relations stored in a semantic graph.

4 Experimental Setup

4.1 TextWorld Interactive Environments

We compared Ariadne agent with alternative methods in a series of text based games involving spatial navigation, object collection and tool manipulation. All these games can be considered Partially Observable MDPs (POMDPs). Such games have long been benchmarks for researching agents capable of effectively remembering information and establishing longterm dependencies [Parisotto et al., 2020; Pleines et al., 2022; Sorokin et al., 2022].

Treasure Hunting. The primary objective is to retrieve the hidden treasure, with a series of rooms providing keys and clues leading to the final goal. The basic variation has 12

contains

Figure 2: AriGraph world model and Ariadne cognitive architecture. (A) AriGraph learns episodic and semantic knowledge during interaction with unknown environment. At each time step t new episodic vertex (containing full textual observation ot) is added to the episodic memory. Then LLMmodel parses observation ot to extract relevant relationships in a form of triplets (object1, relation, object2). These triplets are used to update semantic memory graph. The connection between episodic and semantic memory occurs through episodic edges that link each episodic vertex with all triplets extracted from respective observation. (B) Ariadne agent explores the environment and accomplishes tasks with AriGraph. User sets goal to the agent. Working memory is populated with recent history of observations and actions, relevant semantic and episodic knowledge retrieved from the AirGraph world model. Planing LLM module uses content of working memory to generate new or update existing plan. Results of planning are stored back in working memory. Finally, a ReAct-based module reads memory content and selects one of possible actions to be executed in the environment. Every observation triggers learning that updates agent's world model.

rooms and 4 keys, hard one has 16 rooms and 5 keys and hardest contains 36 rooms, 7 keys and additional distract items in every room.

is south of Cleaning. The goal is to clean a house by identifying and returning misplaced items to their correct locations. Environment consists of 9 rooms (kitchen, pool, etc.) and contains 11 misplaced items (among many other items). To solve the problem, the agent needs to memorize the location of rooms and objects, as well as reason about objects placement.

Cooking. The goal is to prepare and consume a meal by following a recipe, selecting the correct ingredients, and using appropriate tools, while navigating in multi-room house. Basic difficulty task features 9 locations and 3 ingredients and hard task features 12 locations and 4 ingredients, while hardest task also features closed doors and inventory management.

For baselines we used Ariadne's planning and decision making module with one of the following types of memory instead of AriGraph model: full history of observations and actions, iterative summarization, RAG, RAG with Reflexion [Shinn et al., 2023], and Simulacra - memory implementation from [Park et al., 2023].

Full history involves retaining a complete record of all observations and actions to inform decision-making at every step. Summarization, as an alternative to storing the full history, focuses on retaining only the necessary information while discarding the rest. The standard RAG baseline retrieves top-k memories based on their similarity score to the current observation and plan. Simulacra features a scoring mechanism that integrates recency, importance, and relevance, alongside reflections on the extracted memories. The Reflexion baseline differs from other methods in its approach, as it operates over multiple trials. After failing a trial, the agent reflects on its trajectories to document information that may assist in solving the task in subsequent trials. We used the gpt-4-0125-preview as LLM backbone for AriGraph and other LLM-based baselines.

Additionally, we tested our architecture on a variation of the cooking test from [Adhikari et al., 2021] to compare it with RL baselines. These tasks have 4 levels of difficulty, however, they are significantly simpler than our main tasks, having fewer locations, ingredients, and required actions (Appendix F).

For RL baselines, we collect the best results reported by [Adhikari et al., 2021; Tuli et al., 2022; Basu et al., 2024] for the GATA, LTL-GATA, and EXPLORER architectures on the Cooking task with four difficulties levels from [Adhikari et al., 2021].

To estimate human performance in the same games, we developed a graphical user interface , allowing volunteers to play basic versions of the Treasure Hunt, The Cleaning, and the Cooking. After collecting the data, we excluded sessions where the game was not completed.

4.2 NetHack Environment

NetHack [K¨uttler et al., 2020] is a classic roguelike adventure game featuring procedurally generated multi-level dungeon (see Figure 13 in Appendix for a dungeon level example). It poses significant challenges for both LLM-based and RLbased approaches, requiring complex exploration, resource management, and strategic planning.

We based our experiments on NetPlay [Jeurissen et al., 2024] agent, which demonstrates state-of-the-art performance among LLM agents that do not rely on finetuning or RL. In NetPlay agent receives textual observations containing all information about current explored dungeon level. These observations (Level obs) effectively function as handcrafted memory oracle for the agent.

To evaluate our Ariadne agent, we restricted textual observations to agent's current room or corridor (Room Obs), testing whether AriGraph world model could compensate for this restriction by remembering all relevant level information.

We compare three agents. The first is NetPlay [Room obs] with restricted textual observations, the second is our Ariadne [Room obs] agent that receives Room Obs and updates AriGraph, and the last is NetPlay [Level obs] with access to information about explored level.

4.3 Multi-hop Q&A

Although our memory architecture was originally designed for an agent interacting with the environment, we evaluated its performance on standard multi-hop Q&A benchmarks - Musique [Trivedi et al., 2022] and HotpotQA [Yang et al., 2018] to show its robustness and efficiency in more standard retrieval tasks. We made slight adjustments to the promts and replaced Contriever model with BGE-M3[Chen et al., 2024], as it is a better fit for general text encoding. We used 200 random samples from both datasets similar to [Li et al., 2024a]. We compared the performance of our approach against Graphreader [Li et al., 2024a], ReadAgent [Lee et al., 2024], HOLMES [Panda et al., 2024], GraphRAG [Edge et al., 2024] and RAG baselines provided in [Li et al., 2024a].

5 Results

5.1 TextWorld

Every LLM based agent had five attempts to solve each game. The normalized score of one means that an agent completed the game, and score less than one represents intermediate progress. Results on text-based games are shown on the Figure 3 (for dynamics see Appendix G). We estimate performance as average of three best runs. Ariadne successfully remembers and uses information about state of the world for all three tasks. Baseline agents are unable to solve the Treasure Hunt, and fail to find even second key in the Treasure Hunt Hardest. On the other hand, Ariadne successfully solves the Treasure Hunt in about fifty steps, maintains robust performance in the Treasure Hunt Hard, and is able to complete the Treasure Hunt Hardest with more then double amount of rooms compared to Hard version, additional keys and distractors (see Appendix G).

Compared to the Treasure Hunt, the Cleaning game it is more important to properly filter outdated information about object locations, than not to lose any information. This is evident from the reduced usefulness of Episodic Memory in Ariadne agent and Full history baseline, since both memory modules focus on retaining long-term information. Overall Ariadne notably outperforms alternatives in this game. Moreover, Ariadne also outperforms Reflexion, which has additional information between episodes [Shinn et al., 2023]. This baseline shows markable performance growth (in comparison wuth RAG) at the second try, but degrades with following tries.

The Cooking game has the highest difficulty, because any error at intermediate step prevents completion of the whole game. All baseline agents (except Reflexion 2-shot with obvious advantage over other methods) fail to complete cooking tasks due to insufficient or misused information. In this game, episodic memory is particularly important, allowing the agent to recall useful observations such as the content of the recipe or cooking instructions. For token usage of every method see Table 3, Appendix D.

Comparison with RL baselines on variation of the Cooking task is shown in Figure 4. We run Ariadne and GPT-4 with Full history on 4 difficulty levels from the cooking benchmark

[Adhikari et al., 2021]. Ariadne shows superior performance to RL-agents on all 4 levels, especially harder ones. GPT-4 agent with Full history solves only first two levels which is consistent with previous result as the Cooking from Figure 3.A is harder than level 4.

Human evaluation. Comparison with respect to human players is shown in Figure 3.C. All Humans is the average score of all valid (completed) human trials. Human Top-3 is the average score of three best plays for each task. Ariadne outperforms average human player from our sample on all tasks, and scores similarly to the best human plays in the Cooking and the Treasure Hunt, but underperforms in the Cleaning.

Graph quality. We measured AriGraph's growth rate and update rate during gameplay (see Figure 5). The graph actively grows during the exploration phase and flattens once the agent becomes familiar with the environment. We argue that this indicates that agent can generalize to long interactions with the environment despite constant updates to the semantic graph. Additional results in Appendix C demonstrate that the growth rate of the graph decreases with the increase in quality of LLM backbone.

Overal results demonstrate clear advantage of Ariadne agent over LLM based and RL baselines. Semantic memory enables the Ariadne Agent to build and update knowledge about the current state of the POMDP environment, which is crucial for navigation, exploration and capturing relevant details in interactive environments. On the other hand, episodic memory assists the agent in retrieving detailed long-term information that may not be captured in semantic memory, as demonstrated by the results in the Cooking task.

5.2 NetHack

The results are presented in Table 1. Scores column shows average game score across 3 runs, Levels column shows average number of dungeon levels completed by an agent. GPT4o was used for all agents. Underscoring the importance of memory in this task, NetPlay [Level obs] with access to memory oracle achieved the highest scores, while NetPlay [Room obs] with only current room observations performed the worst. Ariadne [Room obs] successfully utilized AriGraph word model, achieving performance comparable to the baseline with memory oracle.

5.3 Multi-hop Q&A

We compared AriGraph with the latest LLM-based approaches that employ knowledge graph construction and retrieval techniques for question answering over documents (Table 2). Our memory architecture, adapted from the Ariadne TextWorld agent, utilizing both GPT-4 and GPT-4omini outperformed baseline methods like ReadAgent (GPT4), GPT-4 RAG, GPT-4 full context and GraphReader (GPT4). GraphRAG served as a strong GPT-4o-mini baseline, due to its extremely hight costs. ArigGraph (GPT-4o-mini) showed weaker performance on Musique, but outperformed GraphRAG on HotpotQA. Notably, our approach is more then 10x cheaper in comparison to GraphRAG (Table 3, Appendix D).

The best performance using GPT-4 was achieved by HOLMES, but AriGraph (GPT-4) exhibited comparable results. Notably, all baseline methods were specifically designed for Q&A tasks, incorporating task-specific prompt tuning and additional architectural enhancements. Both GraphRAG and HOLMES employ hyper-relations in their graphs to connect source data with extracted entities, similar to our method. However, these approaches lack mechanisms for updates in dynamic environments, a key advantage of AriGraph.

6 RelatedWork

Voyager [Wang et al., 2023a], Ghost in the Minecraft [Zhu et al., 2023] and Jarvis-1 [Wang et al., 2023b] are advanced, open-ended LLM agents that show significantly better performance in Minecraft compared to earlier techniques. These agents feature memory capabilities through a library of learned skills, summaries of successful actions, and episodic memory with plans for successful task execution. However, they fall short in representing knowledge with semantic structure and depend heavily on the LLM's extensive Minecraft knowledge. Generative agents [Park et al., 2023] mimic human behavior in multi-agent environments and were among the pioneers in introducing an advanced memory system for LLM agents. Reflexion [Shinn et al., 2023] and CLIN [Majumder et al., 2023] enables agents to reflect on past trajectories, allowing them to store relevant insights about completed actions in a long-term memory module, but has no structural representation of knowledge and episodic memories. LARP [Yan et al., 2023] utilizes the concepts of episodic and semantic memories but treats them as separate instances and lacks a structural representation of knowledge.

Considerable research is dedicated to leveraging established knowledge graphs for enhancing Q&A [Baek et al., 2023; Li et al., 2024b] systems to address the factual knowledge deficiency observed in LLMs. The latest research demonstrating best performance in Q&A tasks includes Graphreader [Li et al., 2024a], HOLMES [Panda et al., 2024], HippoRAG [Guti´errez et al., 2024], GraphRAG [Edge et al., 2024] which all employ the technique of building knowledge graphs from texts. However, these studies do not address the context of functioning within an interactive environment, nor do they take into account the updates to knowledge graphs prompted by new experiences.

Text-based environments [Cˆot´e et al., 2018; Hausknecht et al., 2019; Shridhar et al., 2021; Wang et al., 2022] were originally designed to evaluate reinforcement learning (RL) agents [Guo et al., 2020; Yao et al., 2020; Ammanabrolu et al., 2020; Ammanabrolu and Hausknecht, 2020; Tuli et al., 2022;

Adhikari et al., 2021]. Multiple experiments have already explored the potential of LLMs in these complex scenarios [Tsai et al., 2023; Tan et al., 2023; Momennejad et al., 2023; Ding et al., 2024a]. However raw LLMs show poor results in these games without proper agentic architecture and memory.

7 Conclusions

In this paper, we introduced AriGraph, a novel knowledge graph world model tailored for LLM agents. AriGraph uniquely integrates semantic and episodic memories from textual observations, providing a structured and dynamic representation of knowledge. We evaluated this approach across a range of interactive text-based games and multi-hop Q&A benchmarks, comparing it against existing memory architectures. To test its capabilities comprehensively, we developed a cognitive architecture called Ariadne, which combines AriGraph with planning and decision-making components.

Our results demonstrate that AriGraph significantly outperforms other memory systems in tasks requiring long-term memory, such as decision-making, planning, and exploration in partially observable environments. The structured knowledge representation provided by AriGraph enables efficient retrieval and reasoning, accelerating learning and task completion. Additionally, AriGraph's scalability was evident as it maintained high performance even when the complexity of tasks increased, involving more objects and locations. In multi-hop Q&A benchmarks, AriGraph exhibited competitive performance, underscoring its robustness and adaptability beyond interactive environments.

While promising, our approach can be further enhanced by incorporating multi-modal observations, procedural memories, and more sophisticated graph search methods.

Acknowledgments

This work was supported by the Ministry of Economic Development of the Russian Federation (code 25-139-66879-10003).