diff --git a/src/nets/__init__py b/src/nets/__init__py new file mode 100644 index 0000000..e69de29 diff --git a/src/nets/deepsets.py b/src/nets/deepsets.py new file mode 100644 index 0000000..35afc3b --- /dev/null +++ b/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"]) diff --git a/src/policies/policy.py b/src/policies/policy.py new file mode 100644 index 0000000..4cc8b77 --- /dev/null +++ b/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 + + + + diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/nets/__init__.py b/tests/nets/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/nets/test_deepsets.py b/tests/nets/test_deepsets.py new file mode 100644 index 0000000..2086044 --- /dev/null +++ b/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]) \ No newline at end of file