Merge pull request #381 from GFNOrg/fix-graph

Don't allow removal of nodes with edges

Author: josephdviviano

src/gfn/gym/graph_building.py

@@ -208,6 +208,10 @@ def backward_step(self, states: GraphStates, actions: GraphActions) -> GraphStat
                 # Update node features
                 graph.x = graph.x[mask]
 
+                # Update edge indices
+                assert torch.all(graph.edge_index != node_idx)
+                graph.edge_index[graph.edge_index > node_idx] -= 1
+
         # Handle ADD_EDGE action
         if torch.any(add_edge_mask):
             add_edge_index = torch.where(add_edge_mask)[0]
@@ -264,8 +268,13 @@ def is_action_valid(
                 )
 
                 if backward:
-                    # For backward actions, we need at least one matching node
-                    if not torch.any(equal_nodes):
+                    # For backward actions, we need one matching node
+                    num_equal_nodes = torch.sum(equal_nodes)
+                    if num_equal_nodes != 1:
+                        return False
+                    # And no edges from/to the node
+                    equal_node_idx = torch.where(equal_nodes)[0][0]
+                    if torch.any(graph.edge_index == equal_node_idx):
                         return False
                 else:
                     # For forward actions, we should not have any matching nodes

src/gfn/states.py

@@ -988,9 +988,11 @@ def backward_masks(self) -> TensorDict:
         )
 
         for i, graph in enumerate(self.data.flat):
-            node_class_masks[i, graph.x.flatten()] = True
-            if graph.x is None:
-                continue
+            node_idxs = torch.arange(len(graph.x.flatten()))
+            has_edge = torch.any(
+                node_idxs[:, None] == graph.edge_index.flatten()[None], dim=1
+            )
+            node_class_masks[i, graph.x.flatten()] = ~has_edge
             ei0, ei1 = get_edge_indices(graph.x.size(0), self.is_directed, self.device)
 
             if graph.edge_index is not None and graph.edge_index.size(1) > 0:

testing/test_states.py

@@ -525,9 +525,9 @@ def test_backward_masks(datas):
 
     # Check action type mask
     assert masks[GraphActions.ACTION_TYPE_KEY].shape == (1, 3)
-    assert masks[GraphActions.ACTION_TYPE_KEY][
+    assert not masks[GraphActions.ACTION_TYPE_KEY][
         0, GraphActionType.ADD_NODE
-    ].item()  # Can remove node
+    ].item()  # Can't remove node as it has an edge
     assert masks[GraphActions.ACTION_TYPE_KEY][
         0, GraphActionType.ADD_EDGE
     ].item()  # Can remove edge
@@ -540,7 +540,14 @@ def test_backward_masks(datas):
     available_nodes = (
         torch.bincount(states.tensor.x.flatten(), minlength=states.num_node_classes) > 0
     )
-    assert torch.all(available_nodes == masks[GraphActions.NODE_CLASS_KEY])
+    nodes_with_edges = (
+        torch.bincount(
+            states.tensor.edge_index.flatten(), minlength=states.num_node_classes
+        )
+        > 0
+    )
+    removable_node = available_nodes & ~nodes_with_edges
+    assert torch.all(removable_node == masks[GraphActions.NODE_CLASS_KEY])
 
     # Check edge_class mask
     assert masks[GraphActions.EDGE_CLASS_KEY].shape == (1, states.num_edge_classes)

tutorials/examples/train_graph_triangle.py

@@ -224,7 +224,7 @@ def main(args: argparse.Namespace) -> None:
     )
     parser.add_argument("--seed", type=int, default=1234, help="Random seed")
     parser.add_argument(
-        "--embedding_dim", type=int, default=64, help="Embedding dim for policy heads"
+        "--embedding_dim", type=int, default=128, help="Embedding dim for policy heads"
     )
     parser.add_argument(
         "--num_conv_layers", type=int, default=1, help="Number of GNN layers"
@@ -235,7 +235,7 @@ def main(args: argparse.Namespace) -> None:
     )
     parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate")
     parser.add_argument(
-        "--lr_Z", type=float, default=1e-1, help="Learning rate for logZ"
+        "--lr_Z", type=float, default=5e-2, help="Learning rate for logZ"
     )
     parser.add_argument(
         "--use_buffer", action="store_true", default=True, help="Use replay buffer"