37 lines
780 B
Python
37 lines
780 B
Python
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)
|
|
|
|
B = 25
|
|
ns = 5
|
|
nv = 7
|
|
npath = 20
|
|
|
|
ego_state = torch.rand(B, ns)
|
|
relative_state = torch.rand(B, nv, ns)
|
|
path_x = torch.rand(B, npath)
|
|
path_y = torch.rand(B, npath)
|
|
|
|
sample = {
|
|
"state": ego_state[0],
|
|
"relative_state": relative_state[0],
|
|
"path_x": path_x[0],
|
|
"path_y": path_y[0],
|
|
}
|
|
module(sample)
|
|
|
|
sample = {
|
|
"state": ego_state,
|
|
"relative_state": relative_state,
|
|
"path_x": path_x,
|
|
"path_y": path_y,
|
|
}
|
|
module(sample)
|