Optimizing Job Allocation using Reinforcement Learning with Graph Neural Networks

TL;DR

The paper addresses the Job Allocation Problem (JAP), which seeks to maximize the number of feasible job-to-resource assignments |A| under selection and conflict constraints on a graph G(P ∪ J, S ∪ C). It formulates JAP as a Markov Decision Process and applies a Graph Neural Network with a Context-Aware Embedding (CAE) module to approximate Q-values for edge selections, enabling policy learning without labeled data. Training employs Double Deep Q-Learning with Prioritized Experience Replay to improve stability and sample efficiency. Empirical results on real-world data (Planny) and synthetic graphs (Erdős–Rényi, Barabási–Albert) show the GNN-based RL method outperforms baselines and generalizes to out-of-distribution instances, highlighting the practical potential of RL+GNN for complex scheduling problems. The approach demonstrates how graph-structured representations and edge-level decision modeling can yield scalable, adaptable solutions for resource allocation tasks.

Abstract

Efficient job allocation in complex scheduling problems poses significant challenges in real-world applications. In this report, we propose a novel approach that leverages the power of Reinforcement Learning (RL) and Graph Neural Networks (GNNs) to tackle the Job Allocation Problem (JAP). The JAP involves allocating a maximum set of jobs to available resources while considering several constraints. Our approach enables learning of adaptive policies through trial-and-error interactions with the environment while exploiting the graph-structured data of the problem. By leveraging RL, we eliminate the need for manual annotation, a major bottleneck in supervised learning approaches. Experimental evaluations on synthetic and real-world data demonstrate the effectiveness and generalizability of our proposed approach, outperforming baseline algorithms and showcasing its potential for optimizing job allocation in complex scheduling problems.

Figures (7)

• Figure 1: Example of the problem instance. We have individuals represented by $p_0, \cdots, p_3 \in P$, and jobs denoted by $j_0, \cdots, j_4 \in J$. The red selection edges from set $S$ connect people to jobs, signifying that a person is qualified to do a job. Additionally, there are directed blue conflict edges in set $C$. These connect jobs, indicating that if a person $p$ is assigned to a job vertex $j$, then that person cannot also be assigned to any $j^\prime \in \mathcal{N}_\text{out}(j)$.
• Figure 2: Example of a transition $(s_t, a_t, r_t, s_{t+1})$ in the MDP. When action $\{p_0, j_0\}$ is picked, its edge and the edge between $p_0$ and $j_1$, which conflicts with $j_0$, are removed from the graph in $s_{t+1}$ (depicted as transparent edges).
• Figure 3: Overview of the model architecture. First, a job allocation graph $G(J \cup P, S \cup C)$ with initial node embeddings $\mu^0 \in \mathbb{R}^{\vert P \vert \times 2}, \nu^0 \in \mathbb{R}^{\vert J \vert \times 2}$ is put through $K$ Context-Aware Embedding modules with parameters $\theta_i \in \Theta$ for $i = 1, \cdots, K$. Afterward, $Q$-values can be predicted by doing an inner product of the corresponding vertex embeddings. For example, for the highlighted edge $\{p_0, j_0\}$, $Q_\theta(G, \{p_0, j_0\}) = \mu^{K^T}_0 \nu^K_0$.
• Figure 4: Overview of the Context-Aware Embedding Module. Given a graph $G(P \cup J, S \cup C)$ and vertex embeddings $\mu \in \mathbb{R}^{\vert P \vert \times d_{in}}, \nu \in \mathbb{R}^{\vert J \vert \times d_{in}}$, the module first splits it into two subgraphs $G(P \cup J, S)$ and $G(J, C)$. Then, these subgraphs are put through their own GAT layers, parameterized by $\theta_0, \theta_1 \in \Theta$. The final embeddings of the job vertices are then computed by combining them symmetrically using the learned linear function $f_{\theta_2}$ with parameters $\theta_2 \in \Theta$.
• Figure 5: Out-of-distribution performance of the algorithms when tweaking the individual parameters of the Erdős–Rényi model.

Theorems & Definitions (3)

• Definition 1: Job Allocation Graph
• Definition 2: Maximum Job Allocation
• Definition 3: Markov Decision Process (MDP)