clone original torch_layers

src/mqt/predictor/ml/gnn.py

@@ -2,7 +2,7 @@
 
 import torch
 import torch.nn as nn
-from torch.nn import functional as F
+from torch.nn import functional
 from torch_geometric.nn import (
     AttentionalAggregation,
     Sequential,
@@ -45,10 +45,10 @@ def __init__(self, **kwargs: object) -> None:
             self.edge_embedding = nn.Embedding(self.num_edge_categories, self.edge_embedding_dim)
 
         if self.node_embedding_dim and self.node_embedding_dim == 1:  # one-hot encoding
-            self.node_embedding = lambda x: F.one_hot(x, num_classes=self.num_node_categories).float()
+            self.node_embedding = lambda x: functional.one_hot(x, num_classes=self.num_node_categories).float()
             self.node_embedding_dim = self.num_node_categories
         if self.edge_embedding_dim and self.edge_embedding_dim == 1 and not self.zx:  # one-hot encoding
-            self.edge_embedding = lambda x: F.one_hot(x, num_classes=self.num_edge_categories).float()
+            self.edge_embedding = lambda x: functional.one_hot(x, num_classes=self.num_edge_categories).float()
             self.edge_embedding_dim = self.num_edge_categories
 
         # hidden dimension accounting for multi-head concatenation

src/mqt/predictor/rl/torch_layers.py

@@ -15,7 +15,7 @@
 
 import torch as th
 import torch.nn as nn
-from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
+from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, NatureCNN, is_image_space
 from torch.nn import functional
 
 logger = logging.getLogger("mqt-predictor")
@@ -44,6 +44,7 @@ def __init__(
         cnn_output_dim: int = 256,
         normalized_image: bool = False,
     ) -> None:
+        # TODO we do not know features-dim here before going over all the items, so put something there. This is dirty!
         super().__init__(observation_space, features_dim=1)
 
         extractors: Dict[str, nn.Module] = {}
@@ -56,6 +57,9 @@ def __init__(
                 total_concat_size += cnn_output_dim
             elif key.startswith("graph"):
                 graph_observation_space.append(subspace)
+            elif is_image_space(subspace, normalized_image=normalized_image):
+                extractors[key] = NatureCNN(subspace, features_dim=cnn_output_dim, normalized_image=normalized_image)
+                total_concat_size += cnn_output_dim
             else:
                 # The observation key is a vector, flatten it if needed
                 extractors[key] = nn.Flatten()