149 lines
4.4 KiB
Python
149 lines
4.4 KiB
Python
import torch
|
|
import gym
|
|
import intersim
|
|
import numpy as np
|
|
import json5
|
|
import os
|
|
opj = os.path.join
|
|
|
|
|
|
from src import InteractionDatasetSingleAgent, metrics
|
|
|
|
def basestr(**kwargs):
|
|
"""
|
|
Return base prefix for all files relating to a certain experiment
|
|
Args:
|
|
kwargs (dict): keyword arguments sent to main training loop
|
|
Returns:
|
|
basestr (str): prefix
|
|
"""
|
|
return 'base_'
|
|
|
|
def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs):
|
|
"""
|
|
Main loop for training and testing different imitation models
|
|
Args:
|
|
train (bool): whether to run train loop
|
|
test (bool): whether to run test loop
|
|
method (str): the method to try for imitation
|
|
loc (int): the location index of the roundabout
|
|
config_file (str): path to config file
|
|
kwargs (dict): remaining kwargs for training loop
|
|
"""
|
|
# get/set seed
|
|
seed = kwargs.get('seed',0)
|
|
torch.manual_seed(seed)
|
|
|
|
# make prefix of output files
|
|
outdir = opj('output',method,'loc%02i'%(loc))
|
|
if not os.path.isdir(outdir):
|
|
os.makedirs(outdir)
|
|
filestr = opj(outdir, basestr(**kwargs))
|
|
|
|
# load config
|
|
if config_path:
|
|
with open(config_path, 'r') as cfg:
|
|
config = json5.load(cfg)
|
|
else:
|
|
raise Exception('No config path specified')
|
|
|
|
# method-based training
|
|
if method=='bc':
|
|
from src import bc
|
|
policy_class = bc.BehaviorCloningPolicy
|
|
train_fn = bc.train
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
# default train / cv / test split datasets
|
|
if train:
|
|
|
|
# make policy, train and test datasets, and send to
|
|
policy = policy_class(config)
|
|
train_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[0])#,1,2])
|
|
cv_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[3])
|
|
train_fn(train_dataset, cv_dataset, policy, filestr, **kwargs)
|
|
|
|
if test:
|
|
|
|
# load policy
|
|
policy = policy_class.load_model(config, filestr)
|
|
policy.eval()
|
|
|
|
# simulate policy
|
|
track = 4
|
|
simulate_policy(policy, loc=loc, track=track, filestr=filestr)
|
|
|
|
# run test metrics
|
|
test_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[track])
|
|
metrics(filestr, test_dataset, policy)
|
|
|
|
|
|
def simulate_policy(policy, loc=0, track=0, filestr=''):
|
|
"""
|
|
Simulate a trained policy
|
|
Args:
|
|
policy: the policy to simulate, which should return action directly
|
|
loc (int): location index to test policy
|
|
track (int): track to test policy
|
|
filestr (str): path prefix to save simulation to
|
|
"""
|
|
# animate from environment
|
|
env = gym.make('intersim:intersim-v0', loc=loc, track=track,
|
|
min_acc=-np.inf, max_acc=np.inf)
|
|
|
|
ob, _ = env.reset()
|
|
env.render()
|
|
done = False
|
|
while not done:
|
|
|
|
# get action
|
|
action = policy(ob)
|
|
|
|
# propagate environment
|
|
ob, r, done, info = env.step(action)
|
|
env.render()
|
|
|
|
env.close(filestr=filestr)
|
|
|
|
def parse_args():
|
|
"""
|
|
Parse arguments to main
|
|
Returns:
|
|
kwargs: dictionary of arguments:
|
|
train (bool): whether to run train loop
|
|
test (bool): whether to run test loop
|
|
method (str): the method to try for imitation
|
|
loc (int): the location index of the roundabout
|
|
config (str): config path
|
|
seed (int): RNG seed
|
|
"""
|
|
import argparse
|
|
parser = argparse.ArgumentParser(description='Save Expert Trajectories')
|
|
parser.add_argument('--loc', default=0, type=int,
|
|
help='location (default 0)')
|
|
parser.add_argument("--train", help="train model",
|
|
action="store_true")
|
|
parser.add_argument("--test", help="test model",
|
|
action="store_true")
|
|
parser.add_argument("--method", help="modeling method",
|
|
choices=['bc', 'gail', 'advil'], default='bc')
|
|
parser.add_argument("--config", help="config file path",
|
|
default='config/networks.json5', type=str)
|
|
parser.add_argument('--seed', default=0, type=int,
|
|
help='seed')
|
|
args = parser.parse_args()
|
|
kwargs = {
|
|
'train':args.train,
|
|
'test':args.test,
|
|
'method':args.method,
|
|
'loc':args.loc,
|
|
'config_path':args.config,
|
|
'seed':args.seed
|
|
}
|
|
return kwargs
|
|
|
|
|
|
if __name__ == '__main__':
|
|
kwargs = parse_args()
|
|
main(**kwargs) |