Pytorch Geometric in RLLib?


I’m trying to train a Graph Network on RLLib. For this purpose I’m using the Pytorch Geometric library.

This library has a specific way to collate samples together in a batch (because each sample has a graph of different size, more info here).

I’m a bit lost on how to integrate Pytorch Geometric with RLLib.

Should I convert observation to tensor graph representation in my custom environment ? Or should I do it on model side ?

Alternatively, can I change the way samples are collated into batch ?

Hey @Astariul , have you taken a look at the Repeated space that RLlib offers? It allows for variable sized observations. There is an example script here that shows how you can use it with an env that outputs variable sized observations:

1 Like

I use torch geometric in ray rllib here: graph-conv-memory-paper/ at master · smorad/graph-conv-memory-paper · GitHub