Implement deepsets and deepsets policy, add tests

src/nets/deepsets.py

@@ -0,0 +1,102 @@
+import torch
+from torch import nn
+
+class DeepSetsModule(nn.Module):
+    def __init__(self, input_dim, phi_hidden_n, phi_hidden_dim, latent_dim, rho_hidden_n, rho_hidden_dim, output_dim):
+        """
+        Args:
+            input_dim (int): input size of one instance of the set; input size of phi
+            phi_hidden_n (int): number of hidden layers in phi
+            phi_hidden_dim (int): size of hidden layers in phi
+            latent_dim (int): output size of phi network, where sum is taken over instances; input size of rho
+            rho_hidden_n (int): number of hidden layers in rho
+            rho_hidden_dim (int): size of hidden layers in rho
+            output_dim (int): output size of rho
+        """
+        super(DeepSetsModule, self).__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.phi = Phi(self.input_dim, phi_hidden_n, phi_hidden_dim, latent_dim)
+        self.rho = Phi(latent_dim, rho_hidden_n, rho_hidden_dim, self.output_dim)
+        self.pooling = torch.sum  # torch.max # torch.mean
+
+    @staticmethod
+    def from_config(config):
+        """
+        Args:
+            config (dict): dictionary with network parameters in the form
+                {
+                    "input_dim": 5,
+                    "phi": {
+                        "hidden_n": 1,
+                        "hidden_dim": 10,
+                    },
+                    "latent_dim": 8,
+                    "rho": {
+                        "hidden_n": 1,
+                        "hidden_dim": 10,
+                    },
+                    "output_dim" : 1,
+                }
+        Returns:
+            m (nn.Module): deep sets module
+        """
+        input_dim = config["input_dim"]
+        phi = config["phi"]
+        latent_dim = config["latent_dim"]
+        rho = config["rho"]
+        output_dim = config["output_dim"]
+        m = DeepSetsModule(input_dim, phi["hidden_n"], phi["hidden_dim"], latent_dim, rho["hidden_n"], rho["hidden_dim"], output_dim)
+        return m
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.tensor): (batch_size, dynamic_size, input_dim)
+        Returns:
+            y (torch.tensor): (batch_size, output_dim)
+        """
+        # use negative dynamic_dim since batch dimensions are inserted at the front
+        dynamic_dim = -2
+        # iterate over dynamic dimension to apply phi to every instance
+        latent = tuple(self.phi(instance) for instance in x.unbind(dynamic_dim))
+        # stack outputs of phi
+        latent = torch.stack(latent, dim=dynamic_dim)
+        # apply pooling function to reduce dynamic dimension
+        latent = self.pooling(latent, dim=dynamic_dim)
+        # apply rho network
+        y = self.rho(latent)
+        return y
+
+
+class Phi(nn.Module):
+    def __init__(self, input_dim, hidden_n, hidden_dim, output_dim):
+        """
+        Fully connected feedforward network with same size for all hidden layers and ReLU activation
+
+        Args:
+            input_dim (int): input dimension
+            hidden_n (int): number of hidden layers
+            hidden_dim (int): hidden layer dimension
+            output_dim (int): output dimension
+        """
+        super(Phi, self).__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.layers = [nn.Linear(self.input_dim, hidden_dim)]
+        for _ in range(hidden_n - 1):
+            self.layers.append(nn.Linear(hidden_dim, hidden_dim))
+        self.layers.append(nn.Linear(hidden_dim, self.output_dim))
+        # self.in_layer = nn.Linear(input_dim, hidden_dim)
+        # self.hidden_layers = [nn.Linear(hidden_dim, hidden_dim) for _ in range(hidden_n - 1)]
+        # self.out_layer = nn.Linear(hidden_dim, output_dim)
+        self.activation = nn.functional.relu
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = self.activation(layer(x))
+        return x
+
+    @staticmethod
+    def from_config(config):
+        return Phi(config["input_dim"], config["hidden_n"], config["hidden_dim"], config["output_dim"])

src/policies/policy.py

@@ -0,0 +1,44 @@
+
+import torch
+from torch import nn
+
+from src.nets.deepsets import DeepSetsModule, Phi
+
+class Policy:
+    pass
+
+class DeepSetsPolicy(Policy, nn.Module):
+    def __init__(self, ego_config, dynamic_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
+            path_config (dict): dictionary for configuring the path network
+            head_config (dict): dictionary for configuring the common head network
+        """
+        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"]
+
+    def forward(self, ego_state, relative_states, path):
+        """
+        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
+        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])
+        x = self.head(x)
+        return x
+
+
+
+

tests/nets/test_deepsets.py

@@ -0,0 +1,66 @@
+import torch
+import random
+from src.nets import deepsets as ds
+
+ds_config = {
+    "input_dim": 5,
+    "phi": {
+        "hidden_n": 1,
+        "hidden_dim": 10,
+    },
+    "latent_dim": 8,
+    "rho": {
+        "hidden_n": 1,
+        "hidden_dim": 10,
+    },
+    "output_dim" : 1,
+}
+
+def test_constructor():
+    m = ds.DeepSetsModule.from_config(ds_config)
+
+def test_phi():
+    input_dim = 5
+    phi = ds.Phi(input_dim, 1, 10, 2)
+    x = torch.rand(7, input_dim)
+    y = phi(x)
+    assert y.shape == torch.Size([7, 2])
+
+    y = phi(torch.rand(input_dim))
+    y = phi(torch.rand(7,7,7,input_dim))
+
+def test_deepsets():
+    m = ds.DeepSetsModule.from_config(ds_config)
+
+    input_dim = ds_config["input_dim"]
+    n_dynamic = random.randint(5, 15)
+    x = torch.rand(n_dynamic, input_dim)
+
+    n_batch = 20
+    x = x.unsqueeze(0).expand(n_batch, n_dynamic, input_dim)
+
+    y = m(x)
+    assert y.shape == torch.Size([n_batch, ds_config["output_dim"]])
+
+    for i in range(n_batch):
+        assert torch.allclose(y[i], y[0])
+
+def test_deepsets_computation():
+    n_dynamic = random.randint(5,15)
+    n_batch = 7
+    input_dim = 5
+    output_dim = 3
+    x = torch.rand(n_dynamic, input_dim)
+    x = x.unsqueeze(0).expand(n_batch, n_dynamic, input_dim)
+    assert x.shape == torch.Size([n_batch, n_dynamic, input_dim])
+
+    phi = torch.nn.Linear(input_dim, output_dim)
+
+    y = torch.stack(tuple(phi(instance) for instance in x.unbind(-2)), dim=-2)
+    assert y.shape == torch.Size([n_batch, n_dynamic, output_dim])
+
+    y = y.sum(dim=-2)
+    assert y.shape == torch.Size([n_batch, output_dim])
+
+    for i in range(n_batch):
+        assert torch.allclose(y[i], y[0])