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

adding tool for visualizing acceleration distributions, and making nframes an arg

Browse Source
This commit is contained in:
Arec
2021-07-29 07:35:21 -07:00
parent afc3719ab9
commit 3b4ef6ffb5
4 changed files with 44 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/main.py
View File

@@ -47,7 +47,7 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
# make policy, train and test datasets, and send to
policy = policy_class(config)
train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0])#,1,2])
train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0,1,2])
cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[3])
train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)
@@ -59,7 +59,7 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
# simulate policy
track = 4
simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=500)
simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=kwargs['nframes'])
# run test metrics
test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[track])
@@ -101,4 +101,4 @@ def simulate_policy(policy, loc=0, track=0, filestr='', nframes=float('inf')):
pbar.update()
env.close(filestr=filestr)
env.close(filestr=filestr+'_sim')