Add test and example config for DeepSetsPolicy

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

View File

View File

@@ -0,0 +1,20 @@
import torch
from src.policies.policy import DeepSetsPolicy
import json5
config_path = "config/networks.json5"
with open(config_path, 'r') as cfg:
config = json5.load(cfg)
def test_deepsets_policy():
module = DeepSetsPolicy(config["ego_state"], config["deepsets"], config["path_encoder"], config["head"])
ns = 5
nv = 7
npath = 20
ego_state = torch.rand(ns)
relative_state = torch.rand(nv, ns)
path = torch.rand(npath, 2)
module(ego_state, relative_state, path)