This repo shows an example how TF-GNN can be incorporated for graph (space) use cases in RLlib.
So far, there was no example/support for GNN models in RLlib. I've introduced a new Graph space (built upon Dict space) and added a simple example which shows how to set up an env with an obs in form of a graph. This dict-like obs stores the current state of a graph and gets converted into a GraphTensor (a composite tensor introduced in TF-GNN library). This GraphTensor then is the input for a custom GNN model built with TF-GNN. TF-GNN allows to build a wide variety of state-of-the-art GNNs, see here.
The env and GNN model in my example are for illustration purposes only (no deeper sense). The point is just to show a way how someone can set up an environment with a graph-like observation and how (custom) GNN models can be created for such cases.
More realistic examples might be small use cases of Shortest Path or Traveling Salesman Problem. I've started working on an example to train an agent with RLlib playing the Shannon switching game which is played by two players on a graph. I hope to add this example to the repo in the future.