Files
InteractionImitation/scratch/arec/intersimple/data/process_all_experts.py

49 lines
1.4 KiB
Python

import tqdm
import expert
import copy
import os
import intersim
from tqdm import tqdm
def process_all_experts(filename='expert.pkl',env_args={}, policy_args={}):
"""
Process all experts in the Interaction Dataset
For now, using NormalizedIntersimpleExpert with NRasterizedIncrementingAgent environment
Args:
filename (str): name for track file
env_args (dict): default environment kwargs
policy_args (dict): default policy kwargs
"""
I, J = len(intersim.LOCATIONS), intersim.MAX_TRACKS
pbar = tqdm(total=I*J)
for loc in range(I):
for track in range(J):
it_env_args = copy.deepcopy(env_args)
it_env_args.update({
'loc':loc,
'track':track,
})
out_folder = os.path.join(intersim.LOCATIONS[loc], 'track%04i'%(track))
if not os.path.isdir(out_folder):
os.makedirs(out_folder)
it_path = os.path.join(out_folder,filename)
expert.demonstrations(
expert='NormalizedIntersimpleExpert',
env='NRasterizedIncrementingAgent',
path=it_path,
env_args=it_env_args,
policy_args=policy_args,
)
pbar.update(1)
pbar.close()
if __name__=='__main__':
import fire
fire.Fire(process_all_experts)