huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add test and example config for DeepSetsPolicy

Browse Source
This commit is contained in:
Johannes Fischer
2021-07-19 18:38:50 +02:00
parent 0e6b1e102e
commit 49e32fe37b
4 changed files with 59 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/policies/policy.py
View File

@@ -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