Add test and example config for DeepSetsPolicy

config/networks.json5

@@ -0,0 +1,34 @@
+{
+    ego_state: {
+        input_dim: 5, // number of state vars
+        hidden_n: 1,
+        hidden_dim: 5,
+        output_dim: 5
+    },
+    deepsets: {
+        input_dim: 5, // number of relative state vars for others
+        phi: {
+            hidden_n: 1,
+            hidden_dim: 20,
+        },
+        latent_dim: 20,
+        rho: {
+            hidden_n: 1,
+            hidden_dim: 10,
+        },
+        output_dim: 10
+    },
+    path_encoder: {
+        input_dim: 40, // 2 * path length for (x,y) coordinates
+        hidden_n: 2,
+        hidden_dim: 20,
+        output_dim: 10,
+    },
+    head: {
+        input_dim: 0, // computed in policy constructor
+        hidden_n: 1,
+        hidden_dim: 50,
+        output_dim: 1, // number of outputs e.g. number of actions, or just one
+        final_activation: 'sigmoid',
+    }
+}

src/policies/policy.py

@@ -16,12 +16,14 @@ class DeepSetsPolicy(Policy, nn.Module):
             path_config (dict): dictionary for configuring the path network
             head_config (dict): dictionary for configuring the common head network
         """
+        super(DeepSetsPolicy, self).__init__()
         self.ego_net = Phi.from_config(ego_config)
         self.deepsets = DeepSetsModule.from_config(dynamic_config)
         self.path_net = Phi.from_config(path_config)
-        self.head = Phi.from_config(path_config)
-        output_dim = self.ego_net.output_dim + self.deepsets.output_dim + self.path_net.output_dim
-        assert output_dim == head_config["input_dim"]
+        cat_dim = self.ego_net.output_dim + self.deepsets.output_dim + self.path_net.output_dim
+        # head has number of concatenated features as input
+        head_config["input_dim"] = cat_dim
+        self.head = Phi.from_config(head_config)
 
     def forward(self, ego_state, relative_states, path):
         """
@@ -38,7 +40,3 @@ class DeepSetsPolicy(Policy, nn.Module):
         x = torch.cat([x_ego, x_relative, x_path])
         x = self.head(x)
         return x
-
-
-
-

tests/policies/test_policies.py

@@ -0,0 +1,20 @@
+import torch
+from src.policies.policy import DeepSetsPolicy
+import json5
+
+config_path = "config/networks.json5"
+with open(config_path, 'r') as cfg:
+    config = json5.load(cfg)
+
+def test_deepsets_policy():
+    module = DeepSetsPolicy(config["ego_state"], config["deepsets"], config["path_encoder"], config["head"])
+    
+    ns = 5 
+    nv = 7
+    npath = 20
+
+    ego_state = torch.rand(ns)
+    relative_state = torch.rand(nv, ns)
+    path = torch.rand(npath, 2)
+    
+    module(ego_state, relative_state, path)