Graph observation space with torch_geometric

[Coluding]

Hi all!

I was wondering if you have had any experience modelling a graph observation space? Since there is not direct TensorSpec for that case, I thought of adjusting the torch_geometric graph to a plain tensor, but that could be a bit inconvenient.

Do you have any tips how to use TensorSpec for torch_geometric graph?

Thank you in advance!

[vmoens]

That sounds amazing, wdyt @matteobettini?

[matteobettini]

Hello,

So in BenchMARL we do use torch_geometric in torchrl, but we build the graph in the network (policy) from the agents' observations. The logic we use to build a batched graph from batched tensors is here if that is interesting.

Having directly a graph as an observation (for example in a single agent setting) is a slightly different problem and I think it is interesting to think about it and come up with a torchrl suggested way / set of tools to enable this.

I'll first discuss about a non-batched graph and then about batching since in torchrl we batch densely and torch_gemoetric batches sparsly

Non-batched graphs

A graph in torch_geometric is an instance of an instance of torch_geometric.data.Data with

• data.x: Node feature matrix with shape [num_nodes, num_node_features]
• data.edge_index: Graph connectivity in COO format with shape [2, num_edges] and type torch.long
• data.edge_attr: Edge feature matrix with shape [num_edges, num_edge_features]
• data.y: Target to train against (may have arbitrary shape), e.g., node-level targets of shape [num_nodes, *] or graph-level targets of shape [1, *]. Probably in RL this is not used
• data.pos: Node position matrix with shape [num_nodes, num_dimensions]

These essentially are just tensors.

Because of this I think a normal composite spec should work.

Batched graphs

Some issues start arising when thinking of batching.

The way torch_geoemtric does batching is that it just considers graphs in a batch as disconnected components of one "mega" graph.

So if you have 2 graphs with num_nodes_1, num_nodes_2 and num_edges_1, num_edges_2, a batch in torch geometric is just a graph with num_nodes_1+num_nodes_2 nodes and num_edges_1 + num_edges_2 edges. It keeps the same structure of a single graph described above with the addition of data.batch which will be a tensor of shape [num_nodes_1+num_nodes_2] with ther first num_nodes_1 elements being 0s (indicating those nodes are from graph 0 and the following num_nodes_2 elsements being 1s.

In torchrl, the way we batch data over time and over vectorized environments is that we stack it densely.

In the case the graphs are always the same (aka same num_edges , num_nodes, etc), then there is no problem as the
data of a single graph can be stacked with that of another graph.

But if the graphs start to be different (either over time, or agents, or environemnts) we have to think carefully.

It seems to me that it is impossible pass around data in the nice efficient torch_geometric batching format as torchrl components and tensordict will expect the leading batch dimensions.

So as with all heterogeneous data problems we keep encountering, there are the usual solutiuons:

• LazyStacks, which are essentially lists of tensors, which is what currenty torchrl uses (eg for dynamically changing data and specs)
• padding (meh, i don't like padding)
• nested tensors ?

[Coluding]

Yeah I get your point. As one of the main characteristics of graphs in most scenarios is that their structure is dynamically changing, I think LazyStacks is the best option here.

Thank you for your sophisticated reply. I will work on an implementation now :)