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

Adapt policy input format

Browse Source
This commit is contained in:
Johannes Fischer
2021-07-20 15:45:15 +02:00
parent 7109afb21f
commit 226a427436
2 changed files with 27 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/policies/policy.py
View File

@@ -8,35 +8,40 @@ class Policy:
pass pass
class DeepSetsPolicy(Policy, nn.Module): 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: Args:
ego_config (dict): dictionary for configuring the ego network 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 path_config (dict): dictionary for configuring the path network
head_config (dict): dictionary for configuring the common head network head_config (dict): dictionary for configuring the common head network
""" """
super(DeepSetsPolicy, self).__init__() super(DeepSetsPolicy, self).__init__()
self.ego_net = Phi.from_config(ego_config) 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) 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 has number of concatenated features as input
head_config["input_dim"] = cat_dim head_config["input_dim"] = cat_dim
self.head = Phi.from_config(head_config) self.head = Phi.from_config(head_config)
def forward(self, ego_state, relative_states, path): def forward(self, sample):
""" """
Args: Args:
ego_state (torch.tensor): (ns,) state of ego vehicle sample (dict): sample dictionary with the following entries:
relative_states (torch.tensor): (nv, ns) relative states of other vehicles (dynamic size) state (torch.tensor): (B, 5) raw state
path (torch.tensor): (path_length, 2) coordinates (x,y) of path 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: Returns:
x (torch.tensor): (head_output_dim,) output of common head network x (torch.tensor): (head_output_dim,) output of common head network
""" """
x_ego = self.ego_net(ego_state) ego = self.ego_net(sample["state"])
x_relative = self.deepsets(relative_states) relative = self.deepsets_net(sample["relative_state"])
x_path = self.path_net(path.flatten()) # cat path_x, path_y to tensor of dim (B, 2*P)
x = torch.cat([x_ego, x_relative, x_path]) 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) x = self.head(x)
return x return x

12
tests/policies/test_policies.py
View File

@@ -15,6 +15,14 @@ def test_deepsets_policy():
ego_state = torch.rand(ns) ego_state = torch.rand(ns)
relative_state = torch.rand(nv, ns) relative_state = torch.rand(nv, ns)
path = torch.rand(npath, 2) path_x = torch.rand(npath)
path_y = torch.rand(npath)
module(ego_state, relative_state, path) sample = {
"state": ego_state,
"relative_state": relative_state,
"path_x": path_x,
"path_y": path_y,
}
module(sample)