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

Merge branch 'main' of github.com:sisl/InteractionImitation

Browse Source
This commit is contained in:
Johannes Fischer
2021-08-04 19:44:04 +02:00
parent 5c40de66fa f9729b0a9d
commit 2224e2cd14
20 changed files with 3894 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/main.py
View File

@@ -47,11 +47,11 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
if train:
# make policy, train and test datasets, and send to
train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0,1,2])
train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['train_tracks'])
# generate transform from train_dataset
transforms = generate_transforms(train_dataset)
policy = policy_class(config, transforms)
cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[3])
cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['cv_tracks'])
train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)
if test:
@@ -61,11 +61,10 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
policy.eval()
# simulate policy
track = 4
simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph'])
simulate_policy(policy, loc=loc, track=kwargs['test_tracks'][0], filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph'])
# run test metrics
test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[track])
test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['test_tracks'])
writer = SummaryWriter(filestr)
info = metrics(filestr, test_dataset, policy)
for k, m in info.items():