[Feature] Improve GNN (#93)

* gnn add edge features * global pooling in gnns * global pooling in gnns * amend * amend
facebookresearch · Jun 11, 2024 · d5f952f · d5f952f
1 parent 6c32503
commit d5f952f
Show file tree

Hide file tree

Showing 5 changed files with 296 additions and 87 deletions.
diff --git a/README.md b/README.md
@@ -263,7 +263,7 @@ agent group. Here is a table of the models implemented in BenchMARL
 | Name                           | Decentralized | Centralized with local inputs | Centralized with global input | 
 |--------------------------------|:-------------:|:-----------------------------:|:-----------------------------:|
 | [MLP](benchmarl/models/mlp.py) |      Yes      |              Yes              |              Yes              |
-| [GNN](benchmarl/models/gnn.py) |      Yes      |              No               |              No               |
+| [GNN](benchmarl/models/gnn.py) |      Yes      |              Yes              |              No               |
 | [CNN](benchmarl/models/cnn.py) |      Yes      |              Yes              |              Yes              |
 
 And the ones that are _work in progress_

diff --git a/benchmarl/conf/model/layers/gnn.yaml b/benchmarl/conf/model/layers/gnn.yaml
@@ -6,3 +6,6 @@ self_loops: False
 gnn_class: torch_geometric.nn.conv.GraphConv
 gnn_kwargs:
   aggr: "add"
+
+position_key: null
+velocity_key: null