From 49e32fe37b1a086b274d6d0805dd83dcc8526587 Mon Sep 17 00:00:00 2001
From: Johannes Fischer <johannes.fischer@kit.edu>
Date: Mon, 19 Jul 2021 18:38:50 +0200
Subject: [PATCH] Add test and example config for DeepSetsPolicy

---
 config/networks.json5           | 34 +++++++++++++++++++++++++++++++++
 src/policies/policy.py          | 12 +++++-------
 tests/policies/__init__.py      |  0
 tests/policies/test_policies.py | 20 +++++++++++++++++++
 4 files changed, 59 insertions(+), 7 deletions(-)
 create mode 100644 config/networks.json5
 create mode 100644 tests/policies/__init__.py
 create mode 100644 tests/policies/test_policies.py

diff --git a/config/networks.json5 b/config/networks.json5
new file mode 100644
index 0000000..ae4ead0
--- /dev/null
+++ b/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',
+    }
+}
\ No newline at end of file
diff --git a/src/policies/policy.py b/src/policies/policy.py
index 4cc8b77..d26803f 100644
--- a/src/policies/policy.py
+++ b/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
-
-
-
-
diff --git a/tests/policies/__init__.py b/tests/policies/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
new file mode 100644
index 0000000..5f942a8
--- /dev/null
+++ b/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)
\ No newline at end of file