From 226a427436bf54fc5a4fc2ffe9aaaeb686028f35 Mon Sep 17 00:00:00 2001
From: Johannes Fischer <johannes.fischer@kit.edu>
Date: Tue, 20 Jul 2021 15:45:15 +0200
Subject: [PATCH] Adapt policy input format

---
 src/policies/policy.py          | 29 +++++++++++++++++------------
 tests/policies/test_policies.py | 12 ++++++++++--
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/src/policies/policy.py b/src/policies/policy.py
index d26803f..311a3f0 100644
--- a/src/policies/policy.py
+++ b/src/policies/policy.py
@@ -8,35 +8,40 @@ class Policy:
     pass
 
 class DeepSetsPolicy(Policy, nn.Module):
-    def __init__(self, ego_config, dynamic_config, path_config, head_config):
+    def __init__(self, ego_config, deepsets_config, path_config, head_config):
         """
         Args:
             ego_config (dict): dictionary for configuring the ego network
-            dynamic_config (dict): dictionary for configuring the dynamic input (deepsets) network
+            deepsets_config (dict): dictionary for configuring the deepsets network
             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.deepsets_net = DeepSetsModule.from_config(deepsets_config)
         self.path_net = Phi.from_config(path_config)
-        cat_dim = self.ego_net.output_dim + self.deepsets.output_dim + self.path_net.output_dim
+        cat_dim = self.ego_net.output_dim + self.deepsets_net.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):
+    def forward(self, sample):
         """
         Args:
-            ego_state (torch.tensor): (ns,) state of ego vehicle
-            relative_states (torch.tensor): (nv, ns) relative states of other vehicles (dynamic size)
-            path (torch.tensor): (path_length, 2) coordinates (x,y) of path
+            sample (dict): sample dictionary with the following entries:
+                state (torch.tensor): (B, 5) raw state
+                relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans)
+                path_x (torch.tensor): (B, P) tensor of P future path x positions
+                path_y (torch.tensor): (B, P) tensor of P future path y positions
+                action (torch.tensor): (B, 1) actions taken from each state
         Returns:
             x (torch.tensor): (head_output_dim,) output of common head network
         """
-        x_ego = self.ego_net(ego_state)
-        x_relative = self.deepsets(relative_states)
-        x_path = self.path_net(path.flatten())
-        x = torch.cat([x_ego, x_relative, x_path])
+        ego = self.ego_net(sample["state"])
+        relative = self.deepsets_net(sample["relative_state"])
+        # cat path_x, path_y to tensor of dim (B, 2*P)
+        path = torch.cat([sample["path_x"], sample["path_y"]], dim=-1)
+        path = self.path_net(path)
+        x = torch.cat([ego, relative, path])
         x = self.head(x)
         return x
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index 5f942a8..7ed5153 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -15,6 +15,14 @@ def test_deepsets_policy():
 
     ego_state = torch.rand(ns)
     relative_state = torch.rand(nv, ns)
-    path = torch.rand(npath, 2)
+    path_x = torch.rand(npath)
+    path_y = torch.rand(npath)
+
+    sample = {
+        "state": ego_state,
+        "relative_state": relative_state,
+        "path_x": path_x,
+        "path_y": path_y,
+    }
     
-    module(ego_state, relative_state, path)
\ No newline at end of file
+    module(sample)
\ No newline at end of file