Graph Reinforcement Learning for Network Control
via Bi-Level Optimization

Within the linear model, our commodity quantities evolve in time as

where $f_{i}^{t}$ denotes the change due to flow of commodities along edges and $e_{i}^{t}$ denotes the change due to exchange between commodities at the same graph node. We refer to this expression as the conservation of flow. We also accrue money as

where $m_{f}^{t},m_{e}^{t}\in\mathbb{R}$ denote the money gained due to flows and exchanges respectively. Our overall problem formulation will typically be to control flows and exchanges so as to maximize money over one or more steps subject to additional constraints such as, e.g., flow limitations through a particular edge.

Flows.    We will denote flows along edge $(i,j)$ with $f_{ij}^{t}(k)$. From these flows, we have

which is the net flow (inflows minus outflows). As discussed, associated with each flow is a cost $m_{ij}^{t}(k)$. Note that given this formulation, the total flow cost for all commodities can be written as $m_{ij}^{t}\cdot f_{ij}^{t}=(m_{ij}^{t})^{\top}f_{ij}^{t}$. Thus, we can write the total flow cost at time $t$ as

Exchanges.    To define our exchange relations and their effect on commodity quantities and costs, we will write the effect that exchanges have on money for each node; we write this as $m_{i}^{t}$. Thus, we have $m_{e}^{t}=\sum_{i\in\mathcal{V}}m_{i}^{t}$. We assume there are $N_{e}(i)$ exchange options at each node $i$. The exchange relation takes the form

where $E_{i}^{t}\in\mathbb{R}^{(N_{c}+1)\times N_{e}(i)}$ is an exchange matrix and $w\in\mathbb{R}^{N_{e}(i)}$ are the weights for each exchange. Each column in this exchange matrix denotes an (exogenous) exchange rate between commodities; for example, for $i$’th column $[-1,1,0.1]^{\top}$, one unit of commodity one is exchanged for one unit of commodity two plus $0.1$ units of money. Thus, the choice of exchange weights $w_{i}^{t}$ uniquely determines exchanges $e_{i}^{t}$ and money change due to exchanges, $m_{e}^{t}$.

Convex constraints.    We may impose additional convex constraints on the problem beyond the conservation of flow we have discussed so far. There are a few common examples that one may use in several applications. A common constraint is the non-negativity of commodity values, which we may express as

Note that this inequality is defined element-wise. We may also limit the flow of all commodities through a particular edge via

where this sum could also be weighted per commodity. These linear constraints are only a limited selection of some common examples and the particular choice of constraints is problem-specific.

The previous subsection presented a linear, deterministic problem formulation that yields a convex optimization problem for the decision variables—the chosen flows and exchange weights. However, the formulation is limited by the assumption of linear, deterministic state transitions (among others), and is thus limited in its ability to represent typical real-world systems (please refer to Appendix A for a more complete treatment). In this paper, we focus on solving the nonlinear problem (reflecting real, highly-general problem statements) via a bi-level optimization approach, wherein the linear problem (which has been shown to be extremely useful in practice) is used as an inner control primitive.

We consider a discounted MDP $\mathcal{M}=(\mathcal{S},\mathcal{A},P,R,\gamma)$. Here, $s^{t}\in\mathcal{S}$ is the state and $a^{t}\in\mathcal{A}$ is the action, both at time $t$. The dynamics, $P:\mathcal{S}\times\mathcal{S}\times\mathcal{A}\to[0,1]$ are probabilistic, with $P(s^{t+1}\mid s^{t},a^{t})$ denoting a conditional distribution over $s^{t+1}$. Finally, we use $R:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$ to denote the reward function and $\gamma\in(0,1]$ the discount factor.

State and state space.    Real-world network control problems are typically partially-observed and many features of the world impact the state evolution. However, a small number of features are typically of primary importance, and the impact of the other partially-observed elements can be modeled as stochastic disturbances. Our formulation requires, at each timestep, the commodity values $x^{t}$. Furthermore, the constraint values are required, such as costs, exchange rates, flow capacities, etc. If the graph topology is time-varying, the connectivity at time $t$ is also critical. More precisely, the state elements that we have discussed so far are either properties of the graph nodes (commodity values) or of the edges (such as flow constraints). This difference is of critical importance in our graph neural network architecture.

Generally, the choice of state elements will depend on the information available to a system designer (what can be measured) and on the particular problem setting. Possible examples of further state elements include forecasts of prices, demand and supply, or constraints at future times.

Action and action space.    As discussed in Section 3, an action is defined as all flows and exchanges, $a^{t}=(f^{t},w^{t})$. In the following subsections, we accurately describe the action parametrization under the bi-level formulation.

Dynamics.    The dynamics of the MDP, $P$, describe the evolution of state elements. We split our discussion into two parts: the dynamics associated with commodity and non-commodity elements.

The commodity dynamics are assumed to be reasonably well-modeled by the conservation of flow (1), subject to the constraints; this forms the basis of the bi-level approach that we describe in the next subsection.

The non-commodity dynamics are assumed to be substantially more complex. For example, buying and selling prices may have a complex dependency on past sales, current demand, current supply (commodity values), as well as random exogenous factors. Thus, we place few assumptions on the evolution of non-commodity dynamics and assume that current values are measurable.

Reward.    We assume that our reward is the total discounted money earned over the problem duration. This results in a stage-wise reward function that corresponds to the money earned in that time period, or $R(s^{t},a^{t})=m_{e}^{t}+m_{f}^{t}.$

The previous subsection presented a general MDP formulation that represents a broad class of relevant network optimization problems. The goal is to find a policy $\tilde{\pi}^{*}\in\tilde{\Pi}$ (where $\tilde{\Pi}$ is the space of realizable Markovian policies) such that $\tilde{\pi}^{*}\in\operatorname*{arg\,max}_{\tilde{\pi}\in\tilde{\Pi}}\mathbb{E}_{\tau}\left[\sum_{t=0}^{\infty}\gamma^{t}R(s^{t},a^{t})\right]$, where $\tau=(s^{0},a^{0},s^{1},a^{1},\ldots)$ denotes the trajectory of states and actions. This formulation requires specifying a distribution over all flow/exchange actions, which may be an extremely large space. We instead consider a bi-level formulation

where we compute $a^{t}$ by replacing a single policy that maps from states to actions (i.e., $s^{t}\rightarrow a^{t}$) with two nested policies, mapping from states to desired next states to actions (i.e., $s^{t}\rightarrow\hat{s}^{t+1}\rightarrow a^{t}$). As a consequence of this formulation, the desired next state $\hat{s}^{t+1}$ acts as an intermediate variable, thus avoiding the direct parametrization of an extremely large action space, e.g., flows over edges in a graph. This desired next state is then used in a linear control problem ($\text{LCP}(\cdot,\cdot)$), which leverages a (slightly modified) one-step version of the linear problem formulation of Section 3 to map from desired next state to action. Thus, the resulting formulation is a bi-level optimization problem, whereby the policy $\tilde{\pi}$ is the composition of the policy $\pi(\hat{s}^{t+1}\mid s^{t})$ and the solution to the linear control problem. Specifically, given a sample of $\hat{s}^{t+1}$ from the stochastic policy, we select flow and exchange actions by solving

where $d(\cdot,\cdot)$ is a convex metric which penalizes deviation from the desired next state. The resultant problem is convex and thus may be easily and inexpensively solved to choose actions $a^{t}$, even for very large problems. Please see Appendix B.2, C for a broader discussion.

As is standard in reinforcement learning, we will aim to solve this problem via learning the policy from data. This may be in the form of online learning (Sutton & Barto, 1998) or via learning from offline data (Levine et al., 2020). There are large bodies of work on both problems, and our presentation will generally aim to be as-agnostic-as-possible to the underlying reinforcement learning algorithm used. Of critical importance is the fact that the majority of reinforcement learning algorithms use likelihood ratio gradient estimation (Williams, 1992), which does not require path-wise backpropagation through the inner problem.

We also note that our formulation assumes access to a model (the linear problem) that is a reasonable approximation of the true dynamics over short horizons. This short-term correspondence is central to our formulation: we exploit exact optimization when it is useful, and otherwise push the impacts of the nonlinearity over time to the learned policy. We assume this model is known in our experiments—which we feel is a reasonable assumption across the problem settings we investigate—but it could be learned from state transitions or as learnable parameters in policy learning.

After having introduced the problem formulation and a general framework to control graph-structured systems from experience, here and in Appendix B.1, we broaden the discussion on specific algorithmic components.

Network architectures.    We argue that graph networks represent a natural choice for network optimization problems because of three main properties. First, permutation invariance. Crucially, non-permutation invariant computations would consider each node ordering as fundamentally different and thus require an exponential number of input/output training examples before being able to generalize. Second, locality of the operator. GNNs typically express a local parametric filter (e.g., convolution operator) which enables the same neural network to be applied to graphs of varying size and connectivity and achieve non-parametric expansibility. This is a property of fundamental importance for many real-world graph control problems, which will be dynamic or frequently re-configured, and it is desirable to be able to use the same policy without re-training. Lastly, alignment with the computations used for network optimization problems. As shown in (Xu et al., 2020), GNNs can better match the structure of many network optimization algorithms and are thus likely to achieve better performance.

Action parametrization.    Let us consider the problem of controlling flows in a network. We are interested in defining a desired next state $\hat{s}^{t+1}$ that is ideally (i) lower dimensional, (ii) able to capture relevant aspects for control, and (iii) as-robust-as-possible to domain shifts. At a high level, we achieve this by avoiding the direct parametrization of per-edge desired flow values and compute per-node desired inflow quantities. Concretely, given the total availability $M$ of commodity units in the graph, we define $\hat{s}^{t+1}=\{\hat{q}_{i}^{t+1}\}_{i\in\mathcal{V}},\sum_{i}\hat{q}_{i}^{t+1}=M$ as a desired per-node number of commodity units. We do so by first determining $\tilde{q}_{i}^{t+1}=\{\tilde{q}_{i}^{t+1}\}_{i\in\mathcal{V}}$, where $\tilde{q}_{i}^{t+1}\in[0,1]$ defines the percentage of currently available commodity units to be moved to node $i$ in time step $t$, and $\sum_{i\in\mathcal{V}}\tilde{q}_{i}^{t+1}=1$. We then use this to compute $\hat{q}_{i}^{t+1}=\lfloor\tilde{q}_{i}^{t+1}\cdot M\rfloor$ as the actual number of commodity units. In practice, we achieve this by defining the intermediate policy as a Dirichlet distribution over nodes, i.e., $\pi(\hat{s}^{t+1}|s^{t})=\tilde{q}^{t+1}\cdot M,\tilde{q}^{t+1}\sim\text{Dir}(\tilde{q}^{t+1}|s^{t})$. Crucially, the representation of the desired next state via $\hat{q}_{i}$ (i) is lower-dimensional as it only acts over nodes in the graph, (ii) uses a meaningful aggregated quantity to control flows, and (iii) is scale-invariant by construction as it acts on ratios opposed to raw commodity quantities. Additionally, for problems that require a generation of commodities (e.g., products in a supply chain), we define the desired next state via the exchange weights introduced in Eq (5), $\hat{s}^{t+1}=\{w_{i}^{t+1}\}_{i\in\mathcal{V}},w_{i}^{t+1}\in\mathbb{N}^{+}$, with $w_{i}^{t+1}$ representing the number of commodity units to generate. In practice, this can be achieved by defining the intermediate policy as a Gaussian distribution over nodes (followed by rounding), i.e., $\pi(\hat{s}^{t+1}|s^{t})=\text{round}(w^{t+1}),w^{t+1}\sim\mathcal{N}(w^{t+1}|s^{t})$.

Let us consider an artificial minimum cost flow problem where the goal is to control commodities from one or more source nodes to one or more sink nodes, in the minimum time possible. We assess the capability of our formulation to handle several practically-relevant situations. Specifically, we do so by comparing different versions of our method against an oracle benchmark to investigate the effect of different neural network architectures. Results in Table 1 and in Appendix D.1.3, show how graph-RL approaches are able to achieve close-to-optimal performance in all proposed scenarios while greatly reducing the computation cost compared to traditional solutions (Figure 2 and Appendix C.2)³³3All methods used the same computational CPU resources, namely a AMD Ryzen Threadripper 2950X (16-Core, 32 Thread, 40M Cache, 3.4 GHz base).. Among all formulations (please refer to Appendix D.1 for additional details), MPNN-RL is clearly the best performing architecture, achieving $86.7\%$ of oracle performance, on average. As discussed in Section 4.3, this highlights the importance of the algorithmic alignment (Xu et al., 2020) between the neural network architecture and the nature of the computations needed to solve the task. Crucially, results show how our formulation is able to operate reliably within a broad set of situations, ranging from scenarios characterized by dynamic travel times (Dyn. travel time), dynamic topologies, i.e., with nodes and edges that can be removed or added during an episode (Dyn. topology), capacitated-networks (Capacity) with different depth (2-hop, 3-hop, 4-hop), and multi-commodity problems (Multi-commodity).

In our first real-world experiment, we aim to optimize the performance of a supply chain inventory system. Specifically, this describes the problem of ordering and shipping product inventory within a network of interconnected warehouses and stores in order to meet customer demand while simultaneously minimizing storage and transportation costs. A supply chain system is naturally expressed via a graph $\mathcal{G}=\{\mathcal{V},\mathcal{E}\}$, where $\mathcal{V}=\mathcal{V}_{S}\cup\mathcal{V}_{W}$ is the set of both store $\mathcal{V}_{S}$ and warehouse $\mathcal{V}_{W}$ nodes, and $\mathcal{E}$ the set of edges connecting stores to warehouses. Demand $d_{i}^{t}$ materializes in stores $i\in\mathcal{V}_{S}$ at each period $t$. If inventory is available at the store, it is used to meet customer demand and sold at a price $p$. Unsatisfied orders are maintained over time and are represented as a negative stock (i.e., backorder). At each time step, the warehouse orders additional units of inventory $w_{i}$ from the manufacturers and stores available ones. As commodities travel across the network, they are delayed by transportation times $t_{ij}$. Both warehouses and storage facilities have limited storage capacities $c_{i}$, such that the current inventory $q_{i}$ cannot exceed it. The system incurs a number of operations-related costs: storage costs $m^{S}_{i}$, production costs $m^{O}_{i}$, backorder costs $m^{B}_{i}$, transportation costs $m^{T}_{ij}$.

SCIM Markov decision process.    To apply the methodologies introduced in Section 4, we formulate the SCIM problem as an MDP characterized by the following elements (please refer to Appendix D.2.2 for a formal definition):

Action space ($\mathcal{A}$): we consider the problem of determining (1) the amount of additional inventory $w_{i}$ to order from manufacturers in all warehouse nodes $i\in\mathcal{V}_{W}$, and (2) the flow $f_{ij}$ of commodities to be shipped from warehouses to stores, such that $\mathbf{a}^{t}=\{w_{i}^{t}\}_{i\in\mathcal{V}_{W}}\cup\{f_{ij}^{t}\}_{(i,j)\in\mathcal{E}}$.

Reward $\left(R(s^{t},a^{t})\right)$: we select the reward function in the MDP as the profit of the inventory manager, computed as the difference between sales revenues and costs.

State space ($\mathcal{S}$): the state space describes the current status of the supply network, via node and edge features. Node features contain information on (i) current inventory, (ii) current and estimated demand, (iii) incoming flow, and (iv) incoming orders. Edge features are characterized by (i) travel time $t_{ij}$, and (ii) transportation cost $m^{T}_{ij}$.

Bi-Level formulation.    In what follows and in Appendix D.2.4, we illustrate a specific instantiation of our framework for the SCIM problem. We define the desired outcome $\hat{s}^{t+1}$ as being characterized by two elements: (i) the desired production in warehouse nodes $\hat{w}_{i}^{t+1},\forall i\in\mathcal{V}_{W}$, and (ii) a desired inventory in store nodes $\hat{q}_{i}^{t+1},\forall i\in\mathcal{V}_{S}$.

The LCP selects flow and production actions to best achieve $\hat{s}^{t+1}$ via distance minimization between desired and actual inventory levels. The LCP is further defined by domain-related constraints, such as ensuring that the inventory in store and warehouse nodes does not exceed storage capacity and that shipped products are non-negative and upper bounded by inventory.

Inventory management via graph control.

For the SCIM problem, we define the domain-driven heuristic as a prototypical S-type (or “order-up-to”) policy, which is generally accepted as an effective heuristic (Van Roy et al., 1997). Appendix D.2 provides further experimental details.

Concretely, we measure overall system performance on three different supply chain networks characterized by increasing complexity. Results in Table 2 show that our framework achieves close-to-optimal performance in all tasks. Specifically, Graph-RL achieves 96.8% (1F2S), 96% (1F3S), and 99.5% (1F10S) of oracle performance. Qualitatively, Figure 3 highlights how Graph-RL learns to control the production and shipping policies to match consumer demand while maintaining low inventory storage. More subtly, Figure 3 shows how policies learned through Graph-RL manage to anticipate demand so that products are promptly available in stores by taking production and shipping time under consideration. Results in Table 2 also show how S-type policies, despite being explicitly fine-tuned for all tasks, are largely inefficient and thus incur unnecessary costs and revenue losses, resulting in a profit gap of approximately $15\%$ compared to Graph-RL, on average.

LCP as inductive bias for network computations.    As a further analysis, we compare with an ablation of our framework, which, as in the majority of literature, is defined as a purely end-to-end RL agent that avoids the LCP and directly maps from environment states to production and shipping actions through either MLPs (Peng et al., 2019; Oroojlooyjadid et al., 2022) or GNNs. Results in Figure 4 clearly highlight how the bi-level formulation exhibits significantly improved sample efficiency and performance compared to its end-to-end counterpart, which is either substantially slower at converging to good-quality solutions or does not converge at all, as in Figure 4 (c). We argue that this behavior is due to two main factors: (1) the bi-level agent operates on a lower-dimensional and well-structured representation via $\hat{s}^{t+1}$, and (2) the bi-level formulation provides an implicit inductive bias towards feasible, high-quality solutions via the definition of the LCP. Together, these two properties define an RL agent that exhibits improved efficiency and performance.