TLDR: Vejde is a new framework that combines graph neural networks and reinforcement learning to create AI policies that can learn to make decisions in complex, structured environments and generalize to new, unseen problems of varying sizes and structures. It represents problem states as factual databases, converts them into graphs, and uses neural message passing to learn effective decision-making strategies, showing performance comparable to instance-specific methods while offering broad applicability.
Artificial intelligence, particularly in the realm of reinforcement learning (RL), has shown remarkable capabilities in solving complex decision-making problems. However, a significant challenge remains: how can these AI agents learn to generalize their knowledge to new, unseen situations that might have different sizes or structures than what they were trained on? This is especially crucial for problems involving richly structured data, such as information stored in relational databases, which are common in many real-world applications like automated network incident response.
A new framework called Vejde, developed by Jakob Nyberg and Pontus Johnson, addresses this challenge by combining data abstraction, graph neural networks (GNNs), and reinforcement learning. Vejde aims to create “inductive policy functions” – decision-making strategies that can adapt and perform well across an entire class of problems, rather than being tied to a single, specific instance.
Understanding Vejde’s Approach
At its core, Vejde represents the state of a decision problem not as a simple list of numbers, but as a database of facts about various entities and their relationships. Imagine a network where you have facts about different computers, users, and connections. Vejde takes this collection of facts and transforms it into a special kind of network called a bipartite graph. In this graph, there are two types of nodes: one for the facts themselves and another for the objects (entities) involved in those facts. Edges connect facts to the objects they describe, with each edge also noting the position of the object within the fact.
Once the state is represented as a graph, Vejde employs graph neural networks, specifically a technique known as neural message passing or “graph color refinement.” This process involves nodes in the graph iteratively exchanging information (or “messages”) with their neighbors. Through these exchanges, each node develops a rich, high-dimensional vector representation, often called a “color” or “embedding.” These embeddings implicitly capture the complex relationships and properties within the structured state. A key advantage of GNNs is that their parameters are shared across nodes, meaning the network can handle graphs of varying sizes without needing to be redesigned for each new problem.
From these learned object embeddings, Vejde’s policy function then predicts the probabilities of different actions. The framework is designed to handle actions that can also vary in complexity and number of arguments, making it flexible for diverse problem types.
Evaluating Generalization
To test Vejde’s ability to generalize, the researchers used eight different decision problem domains defined in the Relational Dynamic Influence Diagram Language (RDDL). Each domain included ten problem instances with varying initial conditions and numbers of entities. Crucially, these instances were split into training and testing sets. Vejde agents were trained only on the training instances, and their performance was then evaluated on the unseen test instances to see how well they generalized.
Vejde’s performance was compared against two other types of agents: Multi-Layer Perceptron (MLP) agents and Prost. MLP agents are more traditional neural networks that are trained specifically for each problem instance and cannot generalize to new ones. Prost is a powerful online planning algorithm that searches for optimal actions in real-time, often serving as an upper benchmark for performance.
Also Read:
- Navigating Quantum Networks: A Hybrid AI Approach for Reliable Entanglement Routing
- Enhancing Graph Neural Networks with Flexible Subgraph Pattern Learning
Key Findings and Implications
The evaluation yielded promising results. Vejde policies, when trained through reinforcement learning, showed an average test performance on unseen instances that was not significantly different from the instance-specific MLP policies. This is a major achievement, as Vejde can apply a single learned policy to a whole class of problems, unlike MLPs which require retraining for each new instance or variation.
The study also explored training Vejde through “imitation learning,” where it mimicked actions from the high-performing Prost planner. While imitation learning generally led to higher scores for Vejde, the reinforcement learning approach demonstrated its capability to discover effective inductive policies without relying on expert examples.
While Prost often achieved the highest scores, it requires continuous access to a simulator and performs extensive searches during decision-making, which can be computationally expensive and slow. Vejde, on the other hand, learns a policy once and then makes decisions much faster, only needing the simulator during its training phase.
This research highlights the potential of combining graph learning with model-free reinforcement learning to create AI agents that are not only intelligent but also adaptable. By leveraging the inherent structure of relational data, Vejde offers a path towards more robust and generalizable AI solutions for complex, real-world decision problems. The source code for Vejde is publicly available, facilitating its application to new problem domains. You can find more details about this research in the paper: Vejde: A Framework for Inductive Deep Reinforcement Learning Based on Factor Graph Color Refinement.
